I am using 4 GPUs, let’s say device0, device1, device2, and device3. And the model is:

Inside the init function:

sub_model1.to(device0)
sub_model1 = torch.nn.DataParallel(sub_model1, device_ids=[device0, device1])
sub_model2.to(device2)
sub_model2 = torch.nn.DataParallel(sub_model2, device_ids=[device2, device3])


Inside the forward function:

y = sub_model1(x)
#y = y.to(device2) # we dont need this as we can put data in any device when using DataParallel
out = sub_model2(y)

This gives me out of memory problems after running perfectly for some epochs. The error is written like  (RuntimeError: Caught RuntimeError in replica 1 on device 3.).
Am I doing it correctly? In my case, if I don’t use the DataParallel, it works perfectly with 2GPUs (I simply double the batch_size when using this DataParallel setting, 4 GPUs)