There is a bug in ODEFnc(), for the Parameter can not be augmented assignment directly, or else it would have an error "cannot assign ‘torch.cuda.FloatTensor’ as parameter ‘self.w’ (torch.nn.Parameter or None expected)"
self.w.data = (1 + self.beta) * self.w - self.beta * torch.mm(torch.mm(self.w, torch.t(self.w)), self.w)
xw = torch.einsum('ijkl, lm->ijkm', x, w)
d2 = torch.clamp(self.d2, min=0, max=1)
w2 = torch.mm(self.w2 * d2, torch.t(self.w2))
self.w2.data = (1 + self.beta) * self.w2 - self.beta * torch.mm(torch.mm(self.w2, torch.t(self.w2)), self.w2)
There is a bug in ODEFnc(), for the Parameter can not be augmented assignment directly, or else it would have an error "cannot assign ‘torch.cuda.FloatTensor’ as parameter ‘self.w’ (torch.nn.Parameter or None expected)"
it should be :