self.g_out won’t be gathered, as it’s newly initialized in the replica of your original model.
You could try to initialize it in the __init__ with an empty tensor and assign the output of your generator during the forward pass.
Let me know, if that would work for you.