(self, x, t, model, **model_kwargs)
| 229 | self.schedule = schedule |
| 230 | |
| 231 | def drift(self, x, t, model, **model_kwargs): |
| 232 | model_output = model(x, t, **model_kwargs) |
| 233 | assert model_output.shape == x.shape, "Output shape from ODE solver must match input shape" |
| 234 | return model_output |
| 235 | |
| 236 | def score(self, x, t, model, **model_kwargs): |
| 237 | # we only train velocity, hence only need to compute score from velocity |
no outgoing calls
no test coverage detected