refactor vic reg head

AllenCellModeling · Jul 1, 2024 · a62cf88 · a62cf88
1 parent 3c53f9e
commit a62cf88
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/cyto_dl/nn/head/vic_reg.py b/cyto_dl/nn/head/vic_reg.py
@@ -14,9 +14,11 @@ def __init__(self, dimensions=[2048, 8192, 8192, 8192]):
         super().__init__()
         layers = []
         for i in range(len(dimensions) - 1):
-            layers.append(nn.Linear(dimensions[i], dimensions[i + 1]))
-            layers.append(nn.BatchNorm1d(dimensions[i + 1]))
-            layers.append(nn.ReLU(True))
+            layers += [
+                nn.Linear(dimensions[i], dimensions[i + 1]),
+                nn.BatchNorm1d(dimensions[i + 1]),
+                nn.ReLU(True),
+            ]
         layers.append(nn.Linear(dimensions[-2], dimensions[-1], bias=False))
         self.model = nn.Sequential(*layers)