troubleshooting 5

MannLabs · May 16, 2024 · 3e5ace5 · 3e5ace5
1 parent df8eb0d
commit 3e5ace5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/sparcscore/ml/models.py b/src/sparcscore/ml/models.py
@@ -173,7 +173,7 @@ def __init__(self,
 
         super(VGG2_regression, self).__init__()
 
-        self.norm = nn.BatchNorm2d(5)
+        self.norm = nn.BatchNorm2d(in_channels)
 
         self.features = self.make_layers(self.cfgs[cfg], in_channels)
         self.classifier = self.make_layers_MLP(self.cfgs_MLP[cfg_MLP], self.cfgs[cfg], regression=True) # regression is set to True to make the final layer a single output