i write the following code to test ,it’ right, but the last code segment is still wrong, so the error means i change the hook during running?   besides my model is a DataParallel model… so where is the problem…i am thinking.
import torch
from torch import nn
import numpy as np
QP = []
class SimpleConv(nn.Module):
    def __init__(self):
        super(SimpleConv, self).__init__()
        self.layer1 = nn.Linear(1,32)
        self.layer2 = nn.Linear(32,64)

    def forward(self, x):
        global QP
        QP = np.repeat(x,32)
        handle = self.layer1.register_forward_hook(layer1_hook_fn_forward)
        x = self.layer1(x)
        handle.remove()
        x = self.layer2(x)
        return x

def layer1_hook_fn_forward(module, input, output):
    # this hook will change the output
    global QPs
    a = QP.to(output.device)
    assert a.device == output.device, "tensor device not match"
    return output.mul(a)

for i in range(50):
    model = SimpleConv()
    a= torch.tensor([1.0])
    model(a)