How to change the outputs of forward_hook?

1

Hi,
I would like to change the outputs of module forward_hook, given some torch tensors.

For example, I would like to change the outputs of each relu layers in the module, given select_maps.
Besides, each outputs of relu layer is multiplied by the associated select_maps.

I write a coarse code of my thought as follows:

def hook_forward_net(model, select_maps, layers):
    
    select_output_maps = []
    

    def fun(module, inputs, outputs):
        print(module)
        print('\t'+'inputs size: ', inputs[0].size())
        print('\t'+'outputs size: ', outputs[0].size())
        
        outputs = torch.mul(outputs[0].data.cpu(), select_maps)
    

    
    keys0 = list(model._modules.keys())
    
    for key0 in keys0:
        value0 = model._modules.get(key0)
        
        if type(value0) == torch.nn.modules.container.Sequential:
            
            if type(value0._modules.get('2')) == torch.nn.modules.activation.ReLU:
                
                relu = value0._modules.get('2')
                
                hook = relu.register_forward_hook(fun)

Does anyone know how to do that in pytorch?