def forward(self, x):
if args.method in ['dopri5', 'euler', 'midpoint', 'rk4', 'dopri5_old_cheb'] or \
isinstance(args.method, tuple):
out = odesolver(self.odefunc, x,
options=self.options,
rtol=tol_scheduler(args.rtol, _epoch, args.rtol_scheduler),
atol=tol_scheduler(args.atol, _epoch, args.atol_scheduler), )
else:
# self.options.update({'return_inter_points': args.return_inter_points})
out = odesolver(self.odefunc,
x,
options=self.options, )
if args.method in ['dopri5', 'euler', 'midpoint', 'rk4', 'dopri5_old_cheb'] or \
isinstance(args.method, tuple):
out = out[1, ...]
return out
def forward(self, x):
if self.odefunc.augment_dim > 0:
batch_size, channels, height, width = x.shape
aug = torch.zeros(batch_size, self.odefunc.augment_dim,
height, width).to(self.device)
x_aug = torch.cat([x, aug], 1)
else:
x_aug = x
out = odesolver(self.odefunc, x_aug, self.options)
return out
def forward(self, x, reverse=False):
out = odesolver(self.odefunc, x, self.options, reverse=reverse)
return out
def forward(self, x):
out = odesolver(self.odefunc, x, self.options)
return out