Forward pass for the flow model
(self, x, t, **kwargs)
| 180 | return vector_field |
| 181 | |
| 182 | def forward(self, x, t, **kwargs): |
| 183 | """ |
| 184 | Forward pass for the flow model |
| 185 | """ |
| 186 | if t.numel() == 1: |
| 187 | t = t.expand(x.shape[0]) |
| 188 | _pred = self.sample_vt(x, t, **kwargs) |
| 189 | return _pred |
| 190 | |
| 191 | |
| 192 | def training_losses(self, x1: Tensor, x0: Tensor = None, **cond_kwargs): |