def zero_grad(self, set_to_none: bool=False):
    'Reset gradients from model. If reuse_grad_buffers=True, this will raise an error.'
    if (self.use_gradient_averaging and self.grad_averager.reuse_grad_buffers):
        raise ValueError(f'When running {self.__class__.__name__} with reuse_grad_buffers=True, user should never call zero_grad manually. Gradients will be refreshed internally')
    for param_group in self.param_groups:
        for param in param_group['params']:
            if (param.grad is None):
                pass
            elif set_to_none:
                param.grad = None
            else:
                param.grad.zero_()
Reset gradients from model. If reuse_grad_buffers=True, this will raise an error.
def zero_grad(self, set_to_none: bool=False):
    
    if (self.use_gradient_averaging and self.grad_averager.reuse_grad_buffers):
        raise ValueError(f'When running {self.__class__.__name__} with reuse_grad_buffers=True, user should never call zero_grad manually. Gradients will be refreshed internally')
    for param_group in self.param_groups:
        for param in param_group['params']:
            if (param.grad is None):
                pass
            elif set_to_none:
                param.grad = None
            else:
                param.grad.zero_()