Thank you @ptrblck,
Here what I did now,
grad_bank = {}
model = MyModel()
avg_counter = 0

for epoch in epochs:
    for inputs, targets in batches:
        opt.zero_grad()
        outputs = model.batch(inputs)
        loss = criterion(outputs, targets)
        loss.backward()

        for idx, param in enumerate(model.parameters()):
            grad_bank[f"layer_{idx}"] += param.grad.data
            avg_counter += 1

        opt.step()

for key in grad_bank:
    grad_bank[key] = grad_bank[key] / avg_counter


Is it right? Also, How to use autograd.grad method.
Thanks for help.