Hello Andrea, I have tried implementing the solution you have offered but something is bugging me.
In my case, I have two seperate networks as subnetworks and their output form a loss function together. I have initialized those subnetworks in a larger nn.Module child class and included them in the forward.
However, my reason for applying this was that I have a single GPU which seems to be under allocated (4 GiB allocation in a 32 GiB GPU) and I don’t want to allocate another GPU. But it seems like I am having problems at parallelization of my subnetworks. I was hoping that PyTorch internally allows parallelization by building the computation graph but it seems like something is wrong. My training with two subnetworks takes more than twice duration. Have you observed a similar situation?