I want to confirm my understanding of Hooks in PyTorch.
Objective: Changing incoming Neurons to the Convolution layers during inference. For simplicity, assume the neuron values will be halved.
Code (showing hook only):
# Define hooks
class SaveOutput:
    def __init__(self):
        self.outputs = []

    def __call__(self, module, module_in, module_out):
        # halve neuron values
        module_in = module_in * 0.5
        # see the layer outputs
        self.outputs.append(module_out)

    def clear(self):
        self.outputs = []

save_output = SaveOutput()

# register hooks
hook_handles = []

for layer in model.modules():
    # only for the Conv2D
    if isinstance(layer, torch.nn.modules.conv.Conv2d):
        handle = layer.register_forward_hook(save_output)
        hook_handles.append(handle)

Is this okay?
Thank you.