[Python] Retrieve per-class gradients without create_graph=True and retain_graph=True

I am trying to keep track of the per-class gradients of 10k+ classes. Having retain_graph=True allows to access them but uses significant amount of memory. However, I am trying to figure out a way where I could have both create_graph=False and retain_graph=True while having access to per-class gradients. I have attached my sample code below and am able to set create_graph to False. The line to focus on is:

self._scaler.scale(current_loss).backward(create_graph=create_graph, retain_graph=True)

Any help is appreciated. Thank you in advance!

class NativeGrad:
state_dict_key = "amp_scaler"

def __init__(self):
self._scaler = torch.amp.GradScaler('cuda')

def __call__(self, loss, optimizer, clip_grad=None, model=None, create_graph=False, update_grad=True):
grads: list = []
for current_loss in loss:
for p in model.parameters():
if p.grad is not None:
p.grad.data.zero_()

self._scaler.scale(current_loss).backward(create_graph=create_graph, retain_graph=True) # ← ← ← retain_graph should be False here

grad = torch.cat([
p.grad.flatten().clone() if p.grad is not None else torch.zeros_like(p).flatten()
for p in model.parameters()
])
grads.append(grad)

for p in model.parameters():
if p.grad is not None:
p.grad.data.zero_()

grads_tensor = torch.stack(grads, dim=0)

# Manipulate the grads_tensor...


def state_dict(self):
return self._scaler.state_dict()

def load_state_dict(self, state_dict):
self._scaler.load_state_dict(state_dict)

Continue reading...