Merge pull request #210 from wurining/fix-conv3d-bug

fix that cannot correctly set bn layer channels
sinzlab · Mar 7, 2024 · a71eb25 · a71eb25
2 parents a64a9cc + d59eaab
commit a71eb25
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/neuralpredictors/layers/cores/conv3d.py b/neuralpredictors/layers/cores/conv3d.py
@@ -166,7 +166,7 @@ def __init__(
             padding=(0, input_kernel[1] // 2, input_kernel[2] // 2) if self.padding else 0,
         )
 
-        self.add_bn_layer(layer=layer, hidden_channels=hidden_channels[0])
+        self.add_bn_layer(layer=layer, hidden_channels=self.hidden_channels[0])
 
         if layers > 1 or self.final_nonlinearity:
             if hidden_nonlinearities == "adaptive_elu":
@@ -191,7 +191,7 @@ def __init__(
                 padding=(0, self.hidden_kernel[l][1] // 2, self.hidden_kernel[l][2] // 2) if self.padding else 0,
             )
 
-            self.add_bn_layer(layer=layer, hidden_channels=hidden_channels[l + 1])
+            self.add_bn_layer(layer=layer, hidden_channels=self.hidden_channels[l + 1])
 
             if self.final_nonlinearity or l < self.layers:
                 if hidden_nonlinearities == "adaptive_elu":
@@ -386,7 +386,7 @@ def __init__(
             dilation=(self.temporal_dilation, 1, 1),
         )
 
-        self.add_bn_layer(layer=layer, hidden_channels=hidden_channels[0])
+        self.add_bn_layer(layer=layer, hidden_channels=self.hidden_channels[0])
 
         if layers > 1 or final_nonlin:
             if hidden_nonlinearities == "adaptive_elu":
@@ -417,7 +417,7 @@ def __init__(
                 dilation=(self.hidden_temporal_dilation[l], 1, 1),
             )
 
-            self.add_bn_layer(layer=layer, hidden_channels=hidden_channels[l + 1])
+            self.add_bn_layer(layer=layer, hidden_channels=self.hidden_channels[l + 1])
 
             if final_nonlin or l < self.layers:
                 if hidden_nonlinearities == "adaptive_elu":