Print this: torch.cuda.current_device(). The value you get from the print, can you pass it to the ‘device’ argument? Example if the value is 0, can you try as below?
robust_loss_pytorch.adaptive.AdaptiveLossFunction( num_dims = 4, float_dtype=torch.cuda.FloatTensor, device=0)