Amirhj:

self.device = torch.device( “cuda” if torch.cuda.is_available() else “cpu”)


Can you try this instead? This would ensure we always allocate parameters on device 0
self.device = torch.device(
            "cuda:0" if torch.cuda.is_available() else "cpu")