debugging

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