Background:
I’m doing a distributed PPO (basically gathering data from several worker and training on one learner)
Issue:
Data collection works fine but when I train the network with lines below
            self.critic.optimizer.zero_grad()
            batch_states_values = self.critic.forward(batch_states)
            print('crtitic batch_states_values done')
            critic_loss = F.mse_loss(batch_states_values, batch_REFs)
            print('crtitic critic_loss done')
            critic_loss.backward()
            print('crtitic loss backward done')
            self.critic.optimizer.step()
            print('crtitic step done')

And the output shows:
crtitic batch_states_values done
crtitic critic_loss done
crtitic loss backward done

So it appears to be that the program hangs after the loss backward. What could be the cause?
It works fine on my windows workstations but hangs when I run it on a linux machine