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