You need to use hook.
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = nn.Linear(128, 64)
        self.relu1 = nn.Relu()
        self.layer2 = nn.Linear(64, 64)
        self.relu2 = nn.Relu()

        
        self.register_hook = False
        self.hook = {'layer1':[], 'relu1':[], 'layer2':[], 'relu2':[]}
        
    def forward(self, x):
        out_layer1 = self.layer1(x)
        out_relu1 = self.relu1(out_layer1)
        
        out_layer2 = self.layer2(out_relu1)
        out_relu2 = self.relu2(out_layer2)

        if self.register_hook:
            out_layer1.register_hook(lambda grad: self.hook_fn(grad=grad, name='layer1'))
            out_layer1.retain_grad()

            out_relu1.register_hook(lambda grad: self.hook_fn(grad=grad, name='relu1'))
            out_relu1.retain_grad()

            out_layer2.register_hook(lambda grad: self.hook_fn(grad=grad, name='layer2'))
            out_layer2.retain_grad()

            out_relu2.register_hook(lambda grad: self.hook_fn(grad=grad, name='relu2'))
            out_relu2.retain_grad()

        return out_relu2
    
    def hook_fn(self, grad, name):
        self.hook[name].append(grad)

    def reset_hook(self):
        self.hook = {'layer1':[], 'relu1':[], 'layer2':[], 'relu2':[]}