I’m not familiar with the mentioned repository, but by just skimming through the code it seems multiple GPUs won’t be used?
The fit() function points to this line of code 17, which will only use the default device.
You could try to lower the batch size and see, if the model still converges as you wish.
Also, you could try to use torch.utils.checkpoint 10 to trade compute for memory, but I’m not sure how easy it would be to implement it in this particular repository, as some higher level abstractions seem to be used.