I may be running into these problems given how I set my code up. I wanted the code to be adaptable for zero, one, or multiple GPUs, so I have a ModelWrapper class which keeps track of whether it’s being parallelized or not.
Internal to that I have a member model which is the actual nn.Module being parallelized. It extends both nn.Module and an abstract model class (I use an abstract model class so that I can have multiple kinds of model architectures, but the assumptions are that all models in my project have both an encoder and a decoder, and also implement forward as they are modules).
When initializing the ModelWrapper, I first create the model module. This object has attributes for the encoder and decoder (which are objects also extending nn.Module), and these attributes have attributes which are also modules, e.g., an embedding module, and so on. Once I create the model, if I have more than one GPU, I first wrap it in nn.DataParallel, and then put it on the GPU by calling model.cuda().
When I want to use the model during training, e.g., to compute the loss, I just call model(inputs) (do a forward pass), which returns a tensor.
Perhaps the call to nn.DataParallel is not actually distributing the model parameters on the GPU correctly, given how I wrapped everything in classes?
I did verify that all parameters in my model are on the GPU, and during training all three GPUs are being used by the process.