Hello @tom, thank you for the guidance. Now an error is thrown about trying to call the graph a second time when I call loss.backward(). But if I am correct, I just need the gradient values when computing input_grad and the first graph can be dropped.  Do you have further advice on how to solve this without specifying retain_graph=True so I avoid an OOM error?
Below is my current pseudo-code:.
def explanation_loss(model, x, y, l_one):
    logits = model(x)
    ce = F.cross_entropy(logits, target) 
    
    if l_one==0:
        return ce

    sum_log_hat_y = torch.sum(F.log_softmax(logits, dim=1)).mean()
    # same dims as x
    input_grad, = torch.autograd.grad(sum_log_hat_y, x)
    
    mask_matrix = largest_gradient_mask(input_grad)
    mask_matrix = mask_matrix.float()
    
    dot_prod = torch.dot(torch.flatten(mask_matrix), torch.flatten(input_grad))
    
    right_reasons = l_one * (dot_prod) ** 2
    return ce + right_reasons, mask_matrix 


mask_list = []
for batch in clean_train_loader:
    # Mask matrix and loss are calculated in a single pass
    x, target = batch
    x, target = x.to(device).requires_grad_(True), target.to(device)

    optimizer.zero_grad()
    
    loss, A_batch = explanation_loss(model, x, target, 0.01)
    mask_list.append(A_batch)
    
    loss.backward()  # Runtime Error
    optimizer.step()