def _load_averaged_gradients_into_optimizer_(self):
    'If required, load averaged gradients into optimizer; otherwise simply notify grad averager'
    assert self.use_gradient_averaging
    if self.offload_optimizer:
        pass
    else:
        optimized_param_groups = self.state_averager.optimizer.param_groups
        optimized_parameters = [param for group in optimized_param_groups for param in group['params']]
        with torch.no_grad(), self.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.copy_(averaged_grad, non_blocking=True)
    self.grad_averager.notify_used_averaged_gradients()
If required, load averaged gradients into optimizer; otherwise simply notify grad averager
def _load_averaged_gradients_into_optimizer_(self):
    
    assert self.use_gradient_averaging
    if self.offload_optimizer:
        pass
    else:
        optimized_param_groups = self.state_averager.optimizer.param_groups
        optimized_parameters = [param for group in optimized_param_groups for param in group['params']]
        with torch.no_grad(), self.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.copy_(averaged_grad, non_blocking=True)
    self.grad_averager.notify_used_averaged_gradients()