(self, x, t)
| 30 | self.rtol = rtol |
| 31 | |
| 32 | def forward(self, x, t): |
| 33 | self.integration_time = torch.tensor([0, t]).float().type_as(x) |
| 34 | if self.adjoint: |
| 35 | out = torchdiffeq.odeint_adjoint(self.odefunc, x, self.integration_time, rtol=self.rtol, atol=self.atol, |
| 36 | method=self.method, options=dict(step_size=self.step_size, perturb=self.perturb)) |
| 37 | else: |
| 38 | out = torchdiffeq.odeint(self.odefunc, x, self.integration_time, rtol=self.rtol, atol=self.atol, |
| 39 | method=self.method, options=dict(step_size=self.step_size, perturb=self.perturb)) |
| 40 | |
| 41 | return out[-1] |
| 42 | |
| 43 | |
| 44 | class STBlock(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected