code: def adam(grad, state_sum, nodes, lr, device, only_gpu):
      """calculate gradients according to adam"""
      grad_sum = (grad * grad).mean(1)
      if not only_gpu:
          grad_sum = grad_sum.cpu()
      state_sum.index_add_(0, nodes, grad_sum)  # cpu
      std = state_sum[nodes].to(device)  # gpu
      std_values = std.sqrt_().add_(1e-10).unsqueeze(1)
      grad = lr * grad / std_values  # gpu
  
      return grad
  
imports: import random
  import numpy
  import torch
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  from torch.multiprocessing import Queue
  from torch.nn import init
label_desc: Use image node_id 2 for calling a global function with example usage: adam(grad, state_sum, nodes, lr, device, only_gpu) and returns: grad