Join GitHub today
GitHub is home to over 31 million developers working together to host and review code, manage projects, and build software together.Sign up
Move criterion to compute device automatically #455
One issue with this is that we would then force
self.criterion_ = self.criterion(**criterion_params) if isinstance(self.criterion_, torch.nn.Module): self.criterion_ = self.criterion_.to(self.device)
What do you think?