Update branching.py

google · Sep 28, 2023 · 09d7646 · 09d7646
1 parent 4746cc2
commit 09d7646
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
@@ -281,8 +281,8 @@ def kernel_fn(ks: Kernels, **kwargs) -> Kernel:
                   is_input=ks[0].is_input,
                   diagonal_batch=diagonal_batch,
                   diagonal_spatial=diagonal_spatial,
-                  shape1=None,
-                  shape2=None,
+                  shape1=shape1,
+                  shape2=shape2,
                   batch_axis=batch_axis,
                   channel_axis=channel_axis,
                   mask1=None,