(self, x: Tensor, t: Tensor, cfg_scale=1.0, uc_cond=None, cond_key="y", **kwargs)
| 442 | self.sde_sampler = FlowSDE(schedule=self.schedule) |
| 443 | |
| 444 | def forward(self, x: Tensor, t: Tensor, cfg_scale=1.0, uc_cond=None, cond_key="y", **kwargs): |
| 445 | if t.numel() == 1: |
| 446 | t = t.expand(x.size(0)) |
| 447 | # _pred = self.net(x=x, t=t, **kwargs) |
| 448 | _pred = forward_with_cfg(x, t, self.net, cfg_scale=cfg_scale, uc_cond=uc_cond, cond_key=cond_key, **kwargs) |
| 449 | return _pred |
| 450 | |
| 451 | def ode_fn(self, t, x, **kwargs): |
| 452 | return self(x=x, t=t, **kwargs) |
no test coverage detected