(self, batch: dict)
| 542 | return self.log_images(batch, additional_log_kwargs, **kwargs) |
| 543 | |
| 544 | def get_input(self, batch: dict) -> torch.Tensor: |
| 545 | if self.context_parallel_size > 0: |
| 546 | if not is_context_parallel_initialized(): |
| 547 | initialize_context_parallel(self.context_parallel_size) |
| 548 | |
| 549 | batch = batch[self.input_key] |
| 550 | |
| 551 | global_src_rank = get_context_parallel_group_rank() * self.context_parallel_size |
| 552 | torch.distributed.broadcast(batch, src=global_src_rank, group=get_context_parallel_group()) |
| 553 | |
| 554 | batch = _conv_split(batch, dim=2, kernel_size=1) |
| 555 | return batch |
| 556 | |
| 557 | return batch[self.input_key] |
| 558 | |
| 559 | def apply_ckpt(self, ckpt: Union[None, str, dict]): |
| 560 | if ckpt is None: |
nothing calls this directly
no test coverage detected