I believe I’m seeing a certain loss of functionality after upgrading from PyTorch 0.3.1 to 0.4.0.  Specifically I’m trying to use nn.DataParallel to train, on two GPU’s, a model with a parameter that takes up over half the memory of either GPU.  When the DataParallel library code attempts to replicate the model over both GPU’s it broadcasts the parameters to both, and runs out of GPU memory during the broadcast operation.  In 0.3.1 the broadcast operation was implemented in Python, and contained logic to avoid copy operations for which the source and target devices are the same; in 0.4 this operation has been moved into C code, and I’m guessing that this optimization was removed in the process.
If any of the PyTorch developers is reading this, can you verify that my guess is correct?  More important, can anyone suggest a work-around for this problem?  Am I missing some way to use DataParallel (either the module or the functional form) that will work to replicate a model on multiple GPU’s even when it has a parameter that takes up over half the memory of any single GPU?  (As I write this it occurs to me that someone may object that the gradients for such a model won’t fit in GPU memory anyhow, but that isn’t so, thanks to the wonderful feature of sparse tensor gradients.)
Thanks in advance for any light you can shed.