model argument to 2 from 4

deepskies · Mar 18, 2024 · f1ab0d1 · f1ab0d1
1 parent 7da4a5d
commit f1ab0d1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/scripts/models.py b/src/scripts/models.py
@@ -72,7 +72,7 @@ def model_setup_DE(loss_type, DEVICE):#, INIT_LR=0.001):
         Layer = MuVarLayer
         lossFn = loss_bnll
     #opt = torch.optim.Adam(model.parameters(), lr=INIT_LR)
-    model = torch.nn.Sequential(Model(4), Layer())
+    model = torch.nn.Sequential(Model(2), Layer())
     model = model.to(DEVICE)
     return model, lossFn#, opt