What exactly is the behavior of register forward hook with multiple GPUs?
I want to save the outputs of each layer in my model. For now I have this code:
outputs_layers = []
def save_outputs():
    def hook(module, input, output):
        outputs_layers.append(output.data)
        print(len(outputs_layers))
        return None

    return hook

The problem is that, with multiple GPUs, this does not work; each GPU will receive a fraction of the input, so we need to aggregate the results coming from different GPUs.
This can be done easily, for example by making the outputs_layer a dict and concatenating the outputs with the same key. To make this work, though, we would need to be assured that the order in which the GPUs return the values is always the same, and the same order as the inputs.
So, in general, how can we use forward hooks with multiple GPUs?