(self, radius=0.2, gravity=9.8, adjoint=False)
| 13 | |
| 14 | class BouncingBallExample(nn.Module): |
| 15 | def __init__(self, radius=0.2, gravity=9.8, adjoint=False): |
| 16 | super().__init__() |
| 17 | self.gravity = nn.Parameter(torch.as_tensor([gravity])) |
| 18 | self.log_radius = nn.Parameter(torch.log(torch.as_tensor([radius]))) |
| 19 | self.t0 = nn.Parameter(torch.tensor([0.0])) |
| 20 | self.init_pos = nn.Parameter(torch.tensor([10.0])) |
| 21 | self.init_vel = nn.Parameter(torch.tensor([0.0])) |
| 22 | self.absorption = nn.Parameter(torch.tensor([0.2])) |
| 23 | self.odeint = odeint_adjoint if adjoint else odeint |
| 24 | |
| 25 | def forward(self, t, state): |
| 26 | pos, vel, log_radius = state |
nothing calls this directly
no outgoing calls
no test coverage detected