I made the following loss function to attempt to penalise the highest gradients out of the input. Everything runs fine to begin with however I notice that the extra loss function isn’t achieving anything. So I check to see if the gradient changes in the input after I call grad_loss.backward() and it doesn’t change at all.
Loss function:
class LInfLoss:
    def __init__(self, weight=1, tolerance=0.001):
        self.weight = torch.tensor(weight).float()
        self.weight.requires_grad = True
        self.weight = self.weight.to('cuda:0')
        self.tolerance = tolerance

    def __call__(self, gradient: torch.Tensor):
        threshold = gradient.max() * (1 - self.tolerance)
        grads_over_threshold = gradient.ge(threshold).float() * self.weight
        return grads_over_threshold.sum()

Training loop
    for (images, one_hot_labels) in tqdm(batched_train_data):
        # I collect batch size here because the last batch may have a smaller batch_size
        batch_size = images.shape[0]
        images.requires_grad = True
        
        optimizer.zero_grad()
        # as images is not a parameters optimizer.zero_grad() won't reset it's gradient
        if images.grad is not None:
            images.grad = None

        probabilities = model.forward(images)

        # I want to use .backward() twice rather than autograd because I want to accumulate the gradients
        loss = loss_func(probabilities, one_hot_labels)
        loss.backward(create_graph=True)
        grad_loss = grad_loss_func(images.grad)
        grad_loss.backward()

        optimizer.step()

        labels = one_hot_labels.detach().argmax(dim=1)
        predictions = probabilities.detach().argmax(dim=1)
        num_correct = int(predictions.eq(labels).sum())

        train_data_length += batch_size
        train_correct += num_correct
        train_loss += float(loss.detach()) * batch_size

        # To stop memory leaks
        images.grad = None
        del probabilities
        del loss
        del grad_loss
        del labels
        del predictions
        del num_correct

EDIT:
just to clarify grad_loss_func = LInfLoss