Multi-gpu training

[ruinianxu]

Does anyone ever try to train this work via multi-gpu? I've tried single node with multi-gpu and multi-node via slurm. But I always got the issue of RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.. Turn allow_unused on won't address the issue. Thanks for any suggestions.