@ptrblck Thanks for the reply, I was able to get my code running on multiple GPUs!
I pull loss values from multiple layers (that can be on different devices) and add them together before running backward(). Currently, I have to use .to(device) to bring all these loss values to a single device to run backwards. It looks a bit ugly and it’s redundant when I’m using a single GPU/CPU, but I’m not sure if that’s something I could change? DataParallel is for duplicating a model across different GPUs and CPUs, but could I use it only on the loss values so that I don’t have to use to use .to(device)? If I can use it, then would it slow down my code?
I’m also having an issue where GPU:0 always shows at least some usage. In the few tests I did with the same inputs and changing what layers went on which GPUs, I also had at least 850MiB of usage on GPU:0, even when it shouldn’t have been used at all.
Currently I make sure that my inputs and model are CUDA with the following code:
dtype = torch.cuda.FloatTensor

input_image.type(dtype)

model = model.cuda()

Do I need to do something else in order to prevent GPU:0 from always having extra usage?