Hey, thanks for your answer!
I´m trying your suggested solution into the LightningModule’s forward and I´m afraid it is just returning half the data, that is, just the data in one of the two gpus I´m using, any idea of what’s going on and how can I solve it?
Thanks!
PD: I will move the question to github discussions.