Here is a dummy script to recreate the error.
import torch
import torch.nn as nn
import torchvision.datasets as dset
import torchvision.transforms as transforms

from torchvision.transforms import Compose, ToTensor, RandomAffine, RandomApply, Resize
import torch.nn.functional as F

import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square conv kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5x5 image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = LeNet().to(device=device)

## train, test eval:
loss_metric = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), 0.01)

dataset = dset.FashionMNIST(root='/mnt/sda1/Data/', download=True, train=True,
                            transform=transforms.Compose([
                                transforms.Resize(28),
                                transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,)),
                            ]))
valset = dset.FashionMNIST(root='/mnt/sda1/Data/', download=True, train=False,
                            transform=transforms.Compose([
                                transforms.Resize(28),
                                transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,)),
                            ]))
train_loader = torch.utils.data.DataLoader(dataset, batch_size=128,
                                                shuffle=True, num_workers=4)

# testloader = torch.utils.data.DataLoader(valset, batch_size=128,
#                                                 shuffle=False, num_workers=4)
def freeze_weights(model, hooks = []):
    for name, module in model.named_children():
        if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
            gradient_mask = (module.weight>.05).data
            gradient_mask = gradient_mask.float()
            h_ = module.weight.register_hook(lambda grad:print\
            ("Pre hook shape {}: non zero % {}, gradient mask shape {}".format\
            (grad.shape, torch.count_nonzero(grad)/torch.numel(grad), gradient_mask.shape)))
            h = module.weight.register_hook(lambda grad: grad.mul_(gradient_mask))
            h__ = module.weight.register_hook(lambda grad:print\
            ("Post hook shape {}: non zero % {}".format\
            (grad.shape, torch.count_nonzero(grad)/torch.numel(grad))))
            hooks = hooks + [h] + [h_] + [h__]
            print("{} gradient mask shape {} when created".format(name, gradient_mask.shape))
    if len(hooks)==0:
        raise Exception("hook list is empty!") # should be 15
    return hooks
    
for epoch in range(10):
    if epoch>0:
        curr_hooks = freeze_weights(net)
        print("{} hooks created".format(len(curr_hooks)))
    
    for i, (x,y) in enumerate(train_loader):
        optimizer.zero_grad()

        x, y = x.to(device), y.to(device)
            
        outputs = net(x)

        loss = loss_metric(outputs,y)
        
        loss.backward()
        optimizer.step()
    
    if epoch>0:
        for h in curr_hooks: h.remove()
        curr_hooks = []


This is on Python 3.9.1 and torch version 1.7.1+cu110.
From what I understood, the gradient mask of the last layer/hook is being used in all of the hooks. I think I’d have to define a dict() for the gradient masks to keep track of them?