docstr: An asynchronous function that will average gradients
  sent from trainers.
  Args:
      server_rref (RRef): remote reference to the server
      received_batch_number (int): batch number sent by
          the trainer
      param_loc (int): bucket location sent by the trainer
          containing the gradient
      gradient (torch.Tensor or list): tensor sent by the trainer
docstr_tok: ["An","asynchronous","function","that","will","average","gradients","sent","from","trainers",".","Args",":","server_rref","(","RRef",")",":","remote","reference","to","the","server","received_batch_number","(","int",")",":","batch","number","sent","by","the","trainer","param_loc","(","int",")",":","bucket","location","sent","by","the","trainer","containing","the","gradient","gradient","(","torch.Tensor","or","list",")",":","tensor","sent","by","the","trainer"]
code: def average_gradient(
      server_rref, received_batch_number, param_loc, gradient
  ):
      r"""
      An asynchronous function that will average gradients
      sent from trainers.
      Args:
          server_rref (RRef): remote reference to the server
          received_batch_number (int): batch number sent by
              the trainer
          param_loc (int): bucket location sent by the trainer
              containing the gradient
          gradient (torch.Tensor or list): tensor sent by the trainer
      """
      self = server_rref.local_value()
      if type(gradient) is list:
          gradient = sparse_rpc_format_to_tensor(gradient)
      gradient = gradient.cuda(self.rank)
      fut = torch.futures.Future()
      with self.lock:
          if self.batch_number < received_batch_number:
              self.batch_number = received_batch_number
              self.clear_batch_state()
          self.process_gradient(gradient, param_loc)
          if param_loc not in self.futures:
              self.futures[param_loc] = []
          self.futures[param_loc].append(fut)
          if len(self.futures[param_loc]) == self.trainer_count:
              self.record_straggler_end(self.param_key(param_loc))
              param_loc_avg = self.average(param_loc)
              if not self.use_cuda_rpc:
                  param_loc_avg = param_loc_avg.cpu()
              if param_loc_avg.is_sparse:
                  param_loc_avg = sparse_tensor_to_rpc_format(
                      param_loc_avg
                  )
              for cur_fut in self.futures[param_loc]:
                  cur_fut.set_result(param_loc_avg)
              self.record_batch_end(self.param_key(param_loc))
      return fut
  
imports: import functools
  import threading
  import time
  from abc import ABC, abstractmethod
  import torch
  import torch.distributed.rpc
  from metrics.MetricsLogger import MetricsLogger
  from utils import sparse_rpc_format_to_tensor, sparse_tensor_to_rpc_format
label_desc: Use image node_id 7 for calling the AverageParameterServer obj's underlying member method code with example usage: obj.average_gradient(server_rref, received_batch_number, param_loc, gradient) and returns: fut