def _set_zero_grad(self):
    index = 0
    for (i, p) in enumerate(self.model.parameters()):
        if (p.grad is None):
            continue
        layer_size = self.layer_sizes[i]
        p.grad.zero_()
        index += layer_size
def _set_zero_grad(self):
    index = 0
    for (i, p) in enumerate(self.model.parameters()):
        if (p.grad is None):
            continue
        layer_size = self.layer_sizes[i]
        p.grad.zero_()
        index += layer_size