Add correct stride support - Issue #30

f-dangel · Jan 20, 2020 · 368786c · 368786c
1 parent 31518ca
commit 368786c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/backpack/core/derivatives/conv2d.py b/backpack/core/derivatives/conv2d.py
@@ -205,8 +205,10 @@ def weight_jac_t_mat_prod(self, module, g_inp, g_out, mat, sum_batch=True):
         )
 
         grad_weight = grad_weight.view(
-            num_cols, batch, out_channels * in_channels, k_x, k_y
+            num_cols, batch, -1, grad_weight.shape[2], grad_weight.shape[3]
         )
+        grad_weight = grad_weight.narrow(3, 0, k_x).narrow(4, 0, k_y)
+
         if sum_batch is True:
             grad_weight = grad_weight.sum(1)
             batch = 1