(self, x, timesteps)
| 883 | |
| 884 | |
| 885 | def forward(self, x, timesteps): |
| 886 | """ |
| 887 | Apply the model to an input batch. |
| 888 | |
| 889 | :param x: an [N x C x ...] Tensor of inputs. |
| 890 | :param timesteps: a 1-D batch of timesteps. |
| 891 | :return: an [N x K] Tensor of outputs. |
| 892 | """ |
| 893 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) |
| 894 | |
| 895 | results = [] |
| 896 | h = x.type(self.dtype) |
| 897 | for module in self.input_blocks: |
| 898 | h = module(h, emb) |
| 899 | if self.pool.startswith("spatial"): |
| 900 | results.append(h.type(x.dtype).mean(dim=(2, 3))) |
| 901 | h = self.middle_block(h, emb |
| 902 | |
| 903 | if self.pool.startswith("spatial"): |
| 904 | results.append(h.type(x.dtype).mean(dim=(2, 3))) |
| 905 | h = th.cat(results, axis=-1) |
| 906 | return self.out(h) |
| 907 | else: |
| 908 | h = h.type(x.dtype) |
| 909 | return self.out(h) |
nothing calls this directly
no test coverage detected