Apply the block to a Tensor, conditioned on a timestep embedding. :param x: an [N x C x ...] Tensor of features. :param emb: an [N x emb_channels] Tensor of timestep embeddings. :return: an [N x C x ...] Tensor of outputs.
(self, x, emb, s_cond)
| 426 | self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) |
| 427 | |
| 428 | def forward(self, x, emb, s_cond): |
| 429 | """ |
| 430 | Apply the block to a Tensor, conditioned on a timestep embedding. |
| 431 | :param x: an [N x C x ...] Tensor of features. |
| 432 | :param emb: an [N x emb_channels] Tensor of timestep embeddings. |
| 433 | :return: an [N x C x ...] Tensor of outputs. |
| 434 | """ |
| 435 | return checkpoint( |
| 436 | self._forward, (x, emb, s_cond), self.parameters(), self.use_checkpoint |
| 437 | ) |
| 438 | |
| 439 | |
| 440 | def _forward(self, x, emb, s_cond): |
nothing calls this directly
no test coverage detected