[FIX] Flatten after einsum

f-dangel · Apr 28, 2023 · 942af76 · 942af76
1 parent 6cd4466
commit 942af76
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/test_diag_ggn.py b/test/test_diag_ggn.py
@@ -51,7 +51,7 @@ def _einsum_diag_ggn(
  layer.dilation,
  layer.groups,
  )
- return einsum(equation, *operands)
+ return einsum(equation, *operands).flatten(end_dim=1)
 
 
 def _backpack_diag_ggn(