def METHOD_NAME(self):
    with torch.no_grad():
        for key in self.masks.keys():
            module = self.modules[key]
            delattr(module, "mask_name")
            self.masks[key].requires_grad_(False)
            module.forward = partial(torch.nn.Linear.forward, module)