def _make_gradient_averager(self, **kwargs) -> GradientAverager:
    assert hasattr(self, 'state_averager'), 'must initialize state averager first'
    grad_averager = GradientAverager(dht=self.dht, prefix=f'{self.run_id}_grad_averager', parameters=self.state_averager.main_parameters, min_matchmaking_time=self.matchmaking_time, allreduce_timeout=self.allreduce_timeout, shutdown_timeout=self.shutdown_timeout, next_chunk_timeout=self.next_chunk_timeout, client_mode=self.client_mode, auxiliary=self.auxiliary, start=True, **kwargs)
    if self.offload_optimizer:
        optimized_param_groups = self.state_averager.optimizer.param_groups
        optimized_parameters = [param for group in optimized_param_groups for param in group['params']]
        with grad_averager.get_tensors() as averaged_gradients:
            assert (len(averaged_gradients) == len(optimized_parameters))
            for (opt_param, averaged_grad) in zip(optimized_parameters, averaged_gradients):
                opt_param.grad = averaged_grad
    return grad_averager
def _make_gradient_averager(self, **kwargs) -> GradientAverager:
    assert hasattr(self, 'state_averager'), 'must initialize state averager first'
    grad_averager = GradientAverager(dht=self.dht, prefix=f'{self.run_id}_grad_averager', parameters=self.state_averager.main_parameters, min_matchmaking_time=self.matchmaking_time, allreduce_timeout=self.allreduce_timeout, shutdown_timeout=self.shutdown_timeout, next_chunk_timeout=self.next_chunk_timeout, client_mode=self.client_mode, auxiliary=self.auxiliary, start=True, **kwargs)
    if self.offload_optimizer:
        optimized_param_groups = self.state_averager.optimizer.param_groups
        optimized_parameters = [param for group in optimized_param_groups for param in group['params']]
        with grad_averager.get_tensors() as averaged_gradients:
            assert (len(averaged_gradients) == len(optimized_parameters))
            for (opt_param, averaged_grad) in zip(optimized_parameters, averaged_gradients):
                opt_param.grad = averaged_grad
    return grad_averager