debugging

mlcommons · Dec 6, 2023 · cdc7251 · cdc7251
1 parent 1994259
commit cdc7251
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py b/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py
@@ -114,21 +114,21 @@ def __init__(self,
     super().__init__()
 
     if use_layer_norm:
-      norm_layer = nn.LayerNorm
+      norm_layer = nn.LayerNorm(out_chans)
     else:
-      norm_layer = nn.InstanceNorm2d
+      norm_layer = nn.InstanceNorm2d(out_chans)
     if use_tanh:
       activation_fn = nn.Tanh(inplace=True)
     else:
       activation_fn = nn.LeakyReLU(negative_slope=0.2, inplace=True)
     self.conv_layers = nn.Sequential(
         nn.Conv2d(in_chans, out_chans, kernel_size=3, padding=1, bias=False),
-        nn.InstanceNorm2d(out_chans),
-        nn.LeakyReLU(negative_slope=0.2, inplace=True),
+        norm_layer,
+        activations_fn,
         nn.Dropout2d(dropout_rate),
         nn.Conv2d(out_chans, out_chans, kernel_size=3, padding=1, bias=False),
-        nn.InstanceNorm2d(out_chans),
-        nn.LeakyReLU(negative_slope=0.2, inplace=True),
+        norm_layer,
+        activation_fn,
         nn.Dropout2d(dropout_rate),
     )
 
@@ -148,17 +148,17 @@ def __init__(self,
               ):
     super().__init__()
     if use_layer_norm:
-      norm_layer = nn.LayerNorm
+      norm_layer = nn.LayerNorm(out_chans)
     else:
-      norm_layer = nn.InstanceNorm2d
+      norm_layer = nn.InstanceNorm2d(out_chans)
     if use_tanh:
       activation_fn = nn.Tanh(inplace=True)
     else:
       activation_fn = nn.LeakyReLU(negative_slope=0.2, inplace=True)
     self.layers = nn.Sequential(
         nn.ConvTranspose2d(
             in_chans, out_chans, kernel_size=2, stride=2, bias=False),
-        norm_layer(out_chans),
+        norm_layer,
         activation_fn,
     )