Hi @ptrblck,
I was trying out your solution of model sharding. However, I am having trouble doing loss.backward() Since my loss is a sum of the losses from the 2 decoders I am situating on different GPUs. Would you have pointers about a workaround to this issue?
My error trace specifically is:
Train epch 1, 0.00 s - (Done 1 of 73) Traceback (most recent call last):
  File "main_interlingua.py", line 432, in <module>
    loss, norm_e, norm_d = train(args, train_batch, encoder, decoder, decoder2, encoder_optimizer, decoder_optimizer, decoder_optimizer2, criterion)
  File "main_interlingua.py", line 181, in train
    loss.backward()
  File "/usr0/home/spoddar2/anaconda/lib/python2.7/site-packages/torch/autograd/variable.py", line 167, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, retain_variables)
  File "/usr0/home/spoddar2/anaconda/lib/python2.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    variables, grad_variables, retain_graph)
RuntimeError: arguments are located on different GPUs at /opt/conda/conda-bld/pytorch_1512378422383/work/torch/lib/THC/generated/../generic/THCTensorMathPointwise.cu:269