debugging

mlcommons · Dec 6, 2023 · 027574a · 027574a
1 parent 2381c77
commit 027574a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py b/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py
@@ -114,7 +114,7 @@ def __init__(self,
     super().__init__()
 
     if use_layer_norm:
-      norm_layer = nn.LayerNorm()
+      norm_layer = nn.LayerNorm([out_chans, 320, 320])
     else:
       norm_layer = nn.InstanceNorm2d(out_chans)
     if use_tanh:
@@ -148,7 +148,7 @@ def __init__(self,
               ):
     super().__init__()
     if use_layer_norm:
-      norm_layer = nn.LayerNorm()
+      norm_layer = nn.LayerNorm([out_chans, 320, 320)
     else:
       norm_layer = nn.InstanceNorm2d(out_chans)
     if use_tanh: