Hi, I want to use hook function on DataParallel instance, but I have some uncertain points.
I think below code would work well, but I’m not sure about

Do I need to set lock when I access self.target_outputs and self.target_outputs_grad in the hook function?
Is it guaranteed that inputs are scattered to each GPU following original order? For example, if the inputs were [1, 2, 3, 4] and there are 2 GPUs, Do [1, 2] are fed to GPU #1 and [3, 4] are fed to GPU #2?

class Wrapper:

    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.target_outputs = {}
        self.target_outputs_grad = {}

        def forward_hook(_, __, output):
            self.target_outputs[output.device] = output.detach()

        def backward_hook(_, __, grad_output):
            assert len(grad_output) == 1
            self.target_outputs_grad[grad_output[0].device] = grad_output[0].detach()

        self.target_layer.register_forward_hook(forward_hook)
        self.target_layer.register_backward_hook(backward_hook)

Any comments would be appreciated! Thanks!