Hi everyone,
I’m trying to use nn.DataParallel to have a multi-gpu training, but I encountered the following error. I had a look at the various threads, but I wasn’t able to fix the issue:

RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument #2 ‘weight’; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)

I’ve used the following lines before importing torch to limit the visible GPUs (as I do for the training on a single gpu):
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2, 3"
os.getcwd()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

The model it’s used in this way (with weights from another training):
model = ResNet()
model.load_state_dict(torch.load('rot_weights.pt'))
model = nn.DataParallel(model, device_ids=[0, 2, 3]).to(device)
optimizer = optim.Adam(model.parameters(),
                        lr=learning_rate,
                        betas=(0.9, 0.999),
                        eps=1e-08,
                        weight_decay=0,
                        amsgrad=False)

The model is a ResNet-18 from which I want just the feature extraction part, modified for my regression problem:
class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        self.model =  pretrainedmodels.__dict__['resnet18'](pretrained='imagenet')
        self.regression_layer = nn.Sequential(nn.Linear(512, 6))

    def forward(self, x):
        batch_size ,_,_,_ = x.shape #taking out batch_size from input image
        x = self.model.features(x)
        x = torch.nn.functional.adaptive_avg_pool2d(x,1).reshape(batch_size,-1) # then reshaping the batch_size
        x = self.regression_layer(x)
        x = compute_rotation_matrix_from_ortho6d(x.view(batch_size, -1))

        return x
    
    def compute_rotation_matrix_l2_loss(self, gt_rotation_matrix, predict_rotation_matrix):
        loss_function = nn.MSELoss()
        loss = loss_function(predict_rotation_matrix, gt_rotation_matrix)

        return loss

    def compute_rotation_matrix_geodesic_loss(self, gt_rotation_matrix, predict_rotation_matrix):
        theta = compute_geodesic_distance_from_two_matrices(gt_rotation_matrix, predict_rotation_matrix)
        error = theta.mean()

        return error

Any suggestion would be really appreciated!