Gradient accumulation in the inner loop about learn2learn HOT 2 CLOSED

    from torch import autograd 
    from learn2learn.algorithms import maml_update

    def inner_loop(learner, support_dl, fast_lr, adaptation_steps):
       
        all_loss = []
        # initiliaze gradients to zero for all parameters
        accum_grads = tuple(torch.zeros(param.size(), device=device) for param in learner.module.parameters())

        for step in range(adaptation_steps):
            loss = 0.0

            for inputs in support_dl:
                outputs = learner(**inputs)
                accum_step = len(support_dl.dataset) / support_dl.batch_size # number of steps to accumulate gradients 
                curr_loss = outputs.loss / accum_step
                curr_grads = autograd.grad(curr_loss, learner.module.parameters(), create_graph=False)  # False for FOMAML
                accum_grads = tuple(torch.add(accum_grads[i], curr_grads[i]) for i in range(len(curr_grads)))
                loss += curr_loss.item() # accum loss for logging
                del outputs, curr_loss, curr_grads

            maml_update(learner, lr=fast_lr, grads=accum_grads) # update learner parameters with accumulated gradients
            all_loss.append(loss)

      return learner