(self, batch: dict)
| 517 | return self.log_images(batch, additional_log_kwargs, **kwargs) |
| 518 | |
| 519 | def get_input(self, batch: dict) -> torch.Tensor: |
| 520 | if self.context_parallel_size > 0: |
| 521 | if not is_context_parallel_initialized(): |
| 522 | initialize_context_parallel(self.context_parallel_size) |
| 523 | |
| 524 | batch = batch[self.input_key] |
| 525 | |
| 526 | global_src_rank = get_context_parallel_group_rank() * self.context_parallel_size |
| 527 | torch.distributed.broadcast(batch, src=global_src_rank, group=get_context_parallel_group()) |
| 528 | |
| 529 | batch = _conv_split(batch, dim=2, kernel_size=1) |
| 530 | return batch |
| 531 | |
| 532 | return batch[self.input_key] |
| 533 | |
| 534 | def apply_ckpt(self, ckpt: Union[None, str, dict]): |
| 535 | if ckpt is None: |
nothing calls this directly
no test coverage detected