Update branching.py

google · Sep 28, 2023 · 4746cc2 · 4746cc2
1 parent fc11f58
commit 4746cc2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/neural_tangents/_src/stax/branching.py b/neural_tangents/_src/stax/branching.py
@@ -169,8 +169,8 @@ def _mats_prod(nngps, ntks):
         is_input=ks[0].is_input,
         diagonal_batch=ks[0].diagonal_batch,
         diagonal_spatial=ks[0].diagonal_spatial,
-        shape1=None,
-        shape2=None,
+        shape1=ks[0].shape1,
+        shape2=ks[0].shape2,
         batch_axis=ks[0].batch_axis,
         channel_axis=ks[0].channel_axis,
         mask1=None,