def zero_grad()

in improved_diffusion/fp16_util.py [0:0]

• 5 lines of code
• 3 McCabe index (conditional complexity)

def zero_grad(model_params):
    for param in model_params:
        # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group
        if param.grad is not None:
            param.grad.detach_()
            param.grad.zero_()