Unable to use nn.DataParallel with GraphINN

[shakediel]

Hello,

I wish to train a GraphINN model on multiple gpus using pytorchs nn.DataParallel

I keep getting an error saying that the weights and the inputs are not on the same device

I think that the error stems from the lines:

`

    if has_condition:
        mod_out = node.module(mod_in, c=mod_c, rev=rev, jac=jac)
    
    else:
        mod_out = node.module(mod_in, rev=rev, jac=jac)

`

meaning that the operation isnt called from the self.module_list which holds the modules parameters but from the self.node_list itself, thus the weights are kept on the original device and are not transferred to the new device as intended

Does anyone has any possible direction at solving this problem?
Thanks in advance!