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)
| 940 | self.middle_block.apply(convert_module_to_f32) |
| 941 | |
| 942 | def forward(self, x, timesteps): |
| 943 | """ |
| 944 | Apply the model to an input batch. |
| 945 | |
| 946 | :param x: an [N x C x ...] Tensor of inputs. |
| 947 | :param timesteps: a 1-D batch of timesteps. |
| 948 | :return: an [N x K] Tensor of outputs. |
| 949 | """ |
| 950 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) |
| 951 | |
| 952 | results = [] |
| 953 | h = x.type(self.dtype) |
| 954 | for module in self.input_blocks: |
| 955 | h = module(h, emb) |
| 956 | if self.pool.startswith("spatial"): |
| 957 | results.append(h.type(x.dtype).mean(dim=(2, 3))) |
| 958 | h = self.middle_block(h, emb) |
| 959 | if self.pool.startswith("spatial"): |
| 960 | results.append(h.type(x.dtype).mean(dim=(2, 3))) |
| 961 | h = th.cat(results, axis=-1) |
| 962 | return self.out(h) |
| 963 | else: |
| 964 | h = h.type(x.dtype) |
| 965 | return self.out(h) |
| 966 | |
| 967 | |
| 968 | class NLayerDiscriminator(nn.Module): |
nothing calls this directly
no test coverage detected