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.

def step(self, fetch_stats=False):
        """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.
        """
        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)