Runtime error when trying to train on contact prediction task with FP16

[pocuca]

Hello,

Learning that the Contact prediction task raises a memory issue (even with batch size 1) I tried to run it with FP16=True and the following error occurs:

RuntimeError: arguments are located on different GPUs at /opt/conda/conda-bld/pytorch_1579040055865/work/aten/src/THC/generic/THCTensorIndex.cu:400

I tried to train the models with FP16=False on 4 32GB GPUs and was not able to run the contact prediction task with any of them.

Has anyone had this issue before and what would you recommend in this case?

Thank you!