fix

mlcommons · Dec 7, 2023 · 1f95178 · 1f95178
1 parent 8879598
commit 1f95178
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py b/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py
@@ -136,8 +136,8 @@ def __init__(self,
 
     if use_layer_norm:
       size = int(size)
-      norm_layer = LayerNorm
-      normalized_shape = (out_chans, size, size)
+      norm_layer = nn.GroupNorm
+      normalized_shape = (1, out_chans)
     else:
       norm_layer = nn.InstanceNorm2d
       normalized_shape = out_chans
@@ -174,8 +174,8 @@ def __init__(self,
     super().__init__()
     if use_layer_norm:
       size = int(size)
-      norm_layer = LayerNorm
-      normalized_shape = (out_chans, size, size)
+      norm_layer = nn.GroupNorm
+      normalized_shape = (1, out_chans)
     else:
       norm_layer = nn.InstanceNorm2d
       normalized_shape = out_chans