import torch


def running_mean(log, k, v, p=.95):
  """Update running average of a given training metric."""

  if not k in log:
    log[k] = v
  else:
    if torch.isnan(torch.tensor(log[k])):
      log[k] = v
    elif not torch.isnan(torch.tensor(v)):
      log[k] = p * log[k] + (1-p) * v

  return log[k]