code: def step(self, fetch_stats=False):
          
          if self.strategy == "ps":
              return _distributed_sgd_step(
                  self.workers,
                  self.ps_list,
                  write_timeline=False,
                  fetch_stats=fetch_stats)
          else:
              return _simple_sgd_step(self.workers)
docs: Run a single SGD step.
  
          Arguments:
              fetch_stats (bool): Whether to return stats from the step. This can
                  slow down the computation by acting as a global barrier.