(self, x: torch.Tensor, context: torch.Tensor = None, timesteps: int = None)
| 97 | print(f"{self.__class__.__name__} is using checkpointing") |
| 98 | |
| 99 | def forward(self, x: torch.Tensor, context: torch.Tensor = None, timesteps: int = None) -> torch.Tensor: |
| 100 | if self.checkpoint: |
| 101 | return checkpoint(self._forward, x, context, timesteps) |
| 102 | else: |
| 103 | return self._forward(x, context, timesteps=timesteps) |
| 104 | |
| 105 | def _forward(self, x, context=None, timesteps=None): |
| 106 | assert self.timesteps or timesteps |
nothing calls this directly
no test coverage detected