| 360 | self.abondon_cond = abondon_cond |
| 361 | |
| 362 | def forward(self, x: Tensor, vec: Tensor, |
| 363 | pe: Tensor, mask: Tensor = None, |
| 364 | txt_length=None, |
| 365 | uncondi_length=None, |
| 366 | uncondi_pe = None, |
| 367 | mask_uncond = None): |
| 368 | # pad_sequence(tuple(x_list), batch_first=True) |
| 369 | if self.abondon_cond: |
| 370 | x = [ix[:u_l, :] for ix, u_l in zip(x, uncondi_length)] |
| 371 | x = pad_sequence(x, batch_first=True) |
| 372 | if not x.shape[1] == pe.shape[2]: |
| 373 | pe = uncondi_pe |
| 374 | mask = mask_uncond |
| 375 | # print("double stream block", x.shape, pe.shape) |
| 376 | x = super().forward(x, vec, pe, mask, txt_length) |
| 377 | return x |
| 378 | |
| 379 | class SingleStreamBlockC(SingleStreamBlock): |
| 380 | """ |