forward loop of sde
(self, init, model, **model_kwargs)
| 67 | return sampler |
| 68 | |
| 69 | def sample(self, init, model, **model_kwargs): |
| 70 | """forward loop of sde""" |
| 71 | x = init |
| 72 | mean_x = init |
| 73 | samples = [] |
| 74 | sampler = self.__forward_fn() |
| 75 | for ti in self.t[:-1]: |
| 76 | with th.no_grad(): |
| 77 | x, mean_x = sampler(x, mean_x, ti, model, **model_kwargs) |
| 78 | samples.append(x) |
| 79 | |
| 80 | return samples |
| 81 | |
| 82 | |
| 83 | class ode: |