Fix issues with lazy grad initialization (pytorch#912)

Author: apaszke

test/test_nn.py

@@ -357,19 +357,31 @@ def bw_hook(module, grad_input, grad_output):
  self.assertEqual(input.grad.data, expected_grad)
 
  def test_zero_grad(self):
+ i = Variable(torch.randn(2, 5), requires_grad=True)
  module = nn.Linear(5, 5)
  for p in module.parameters():
  p.requires_grad = False
  module.zero_grad()
 
  module.weight.requires_grad = True
- module.weight._grad = Variable(module.weight.data.clone().fill_(1))
+ module.zero_grad()
+ self.assertIsNone(module.weight.grad) # uninitialized grad
+
+ module(i).sum().backward()
+ self.assertIsNotNone(module.weight.grad)
+ self.assertGreater(module.weight.grad.data.abs().sum(), 0)
  module.zero_grad()
  self.assertEqual(module.weight.grad.data, module.weight.data.clone().zero_())
 
  module.bias.requires_grad = True
- module.weight._grad = Variable(module.weight.data.clone().fill_(1))
- module.bias._grad = Variable(module.bias.data.clone().fill_(1))
+ module.zero_grad()
+ self.assertIsNotNone(module.weight.grad)
+ self.assertIsNone(module.bias.grad)
+ module(i).sum().backward()
+ self.assertIsNotNone(module.weight.grad)
+ self.assertIsNotNone(module.bias.grad)
+ self.assertGreater(module.weight.grad.data.abs().sum(), 0)
+ self.assertGreater(module.bias.grad.data.abs().sum(), 0)
  module.zero_grad()
  self.assertEqual(module.weight.grad.data, module.weight.data.clone().zero_())
  self.assertEqual(module.bias.grad.data, module.bias.data.clone().zero_())

torch/nn/modules/module.py

@@ -377,7 +377,7 @@ def eval(self):
  def zero_grad(self):
  """Sets gradients of all model parameters to zero."""
  for p in self.parameters():
- if p.requires_grad:
+ if p.grad is not None:
  p.grad.data.zero_()
 
  def share_memory(self):

torch/nn/utils/clip_grad.py

@@ -11,7 +11,7 @@ def clip_grad_norm(parameters, max_norm, norm_type=2):
  max_norm (float or int): max norm of the gradients
  norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for infinity norm.
  """
- parameters = list(parameters)
+ parameters = list(filter(lambda p: p.grad is not None, parameters))
  max_norm = float(max_norm)
  norm_type = float(norm_type)
  if norm_type == float('inf'):