(self, x, **kwargs)
| 383 | self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) |
| 384 | |
| 385 | def forward(self, x, **kwargs): |
| 386 | # TODO add crossframe attention and use mixed checkpoint |
| 387 | return checkpoint( |
| 388 | self._forward, (x,), self.parameters(), True |
| 389 | ) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!! |
| 390 | # return pt_checkpoint(self._forward, x) # pytorch |
| 391 | |
| 392 | def _forward(self, x): |
| 393 | b, c, *spatial = x.shape |
nothing calls this directly
no test coverage detected