Solved by albanD in post #5 
              
              
                So testing the code and changing the hook to: 
    def __call__(self, grad):
        try:
            grad_clone = grad.clone().detach()

            if self.is_bias:
                grad_clone[self.currently_active] = 0

            else:
                grad_clone[self.currently_active, :] = 0
   …