Fix for non-contiguous grad_output in cuDNN conv

apaszke · apaszke · commit d5b1e0685a39 · 2017-02-09T22:51:50.000+01:00
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -1274,6 +1274,22 @@ def test_inplace_thnn(self):
         output.backward(grad_output)
         self.assertEqual(grad_output, grad_output_clone)
 
+    @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
+    def test_noncontig_conv_grad(self):
+        # FIXME: remove after adding non-contiguous grad tests for all modules
+        module = nn.Conv2d(3, 5, kernel_size=3, padding=1).cuda()
+        input = Variable(torch.randn(2, 3, 10, 10).cuda(), requires_grad=True)
+        output = module(input)
+
+        grad = torch.randn(2, 2, 5, 10, 10).cuda()[:, 1]
+        assert not grad.is_contiguous()
+        output.backward(grad, retain_variables=True)
+        result = output.grad.data.clone()
+        output.grad.data.zero_()
+
+        output.backward(grad.contiguous())
+        self.assertEqual(result, output.grad.data)
+
     def test_pixel_shuffle(self):
         batch_size = random.randint(1, 3)
         upscale_factor = random.randint(2, 5)
diff --git a/torch/nn/_functions/conv.py b/torch/nn/_functions/conv.py
@@ -38,6 +38,7 @@ def forward(self, input, weight, bias=None):
 
     def backward(self, grad_output):
         k = grad_output.dim()
+        grad_output = grad_output.contiguous()
         input, weight, bias = self.saved_tensors
         if k == 3:
             grad_output, input, weight = _view4d(grad_output, input, weight)
