FIX: non-contiguous gradients Closes #31

Former-commit-id: ab06e72
gpauloski · Apr 22, 2021 · 99da3b5 · 99da3b5
1 parent 54179fc
commit 99da3b5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kfac/layers/base.py b/kfac/layers/base.py
@@ -476,9 +476,9 @@ def _get_precondition_gradient_inv(self):
 
     def _set_bias_grad(self, grad):
         """Set bias.grad tensor of module"""
-        self.module.bias.grad = grad
+        self.module.bias.grad = grad.contiguous()
 
     def _set_weight_grad(self, grad):
         """Set weight.grad tensor of module"""
-        self.module.weight.grad = grad
+        self.module.weight.grad = grad.contiguous()