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)
| 984 | self.middle_block.apply(convert_module_to_f32) |
| 985 | |
| 986 | def forward(self, x, timesteps): |
| 987 | """ |
| 988 | Apply the model to an input batch. |
| 989 | :param x: an [N x C x ...] Tensor of inputs. |
| 990 | :param timesteps: a 1-D batch of timesteps. |
| 991 | :return: an [N x K] Tensor of outputs. |
| 992 | """ |
| 993 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) |
| 994 | |
| 995 | results = [] |
| 996 | h = x.type(self.dtype) |
| 997 | for module in self.input_blocks: |
| 998 | h = module(h, emb) |
| 999 | if self.pool.startswith("spatial"): |
| 1000 | results.append(h.type(x.dtype).mean(dim=(2, 3))) |
| 1001 | h = self.middle_block(h, emb) |
| 1002 | if self.pool.startswith("spatial"): |
| 1003 | results.append(h.type(x.dtype).mean(dim=(2, 3))) |
| 1004 | h = th.cat(results, axis=-1) |
| 1005 | return self.out(h) |
| 1006 | else: |
| 1007 | h = h.type(x.dtype) |
| 1008 | return self.out(h) |
nothing calls this directly
no test coverage detected