code: def average_grad(p, mom, dampening=False, grad_avg=None, **kwargs):
      "Keeps track of the avg grads of `p` in `state` with `mom`."
      if grad_avg is None:
          grad_avg = torch.zeros_like(p.grad.data)
      damp = 1 - mom if dampening else 1.0
      grad_avg.mul_(mom).add_(p.grad.data, alpha=damp)
      return {"grad_avg": grad_avg}
  
imports: from __future__ import annotations
  from .torch_basics import *
label_desc: Use image node_id 5 for calling a global function with example usage: average_grad(p, mom, dampening, grad_avg) and returns: dict