(
self,
z: torch.Tensor,
input_cp: bool = False,
output_cp: bool = False,
use_cp: bool = True,
**kwargs,
)
| 592 | return z |
| 593 | |
| 594 | def decode( |
| 595 | self, |
| 596 | z: torch.Tensor, |
| 597 | input_cp: bool = False, |
| 598 | output_cp: bool = False, |
| 599 | use_cp: bool = True, |
| 600 | **kwargs, |
| 601 | ): |
| 602 | if self.cp_size <= 1: |
| 603 | use_cp = False |
| 604 | if self.cp_size > 0 and use_cp and not input_cp: |
| 605 | if not is_context_parallel_initialized: |
| 606 | initialize_context_parallel(self.cp_size) |
| 607 | |
| 608 | global_src_rank = get_context_parallel_group_rank() * self.cp_size |
| 609 | torch.distributed.broadcast(z, src=global_src_rank, group=get_context_parallel_group()) |
| 610 | |
| 611 | z = _conv_split(z, dim=2, kernel_size=1) |
| 612 | |
| 613 | x = super().decode(z, use_cp=use_cp, **kwargs) |
| 614 | |
| 615 | if self.cp_size > 0 and use_cp and not output_cp: |
| 616 | x = _conv_gather(x, dim=2, kernel_size=1) |
| 617 | |
| 618 | return x |
| 619 | |
| 620 | def forward( |
| 621 | self, |
no test coverage detected