corrected n_out bug with optuna

supsi-dacd-isaac · Dec 12, 2023 · 2d5cc09 · 2d5cc09
1 parent f2a5724
commit 2d5cc09
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/pyforecaster/forecasting_models/neural_forecasters.py b/pyforecaster/forecasting_models/neural_forecasters.py
@@ -197,7 +197,6 @@ def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_pat
         if self.load_path is not None:
             self.load(self.load_path)
 
-        self.n_out = self.n_out*2 if self.probabilistic else self.n_out
         self.model = self.set_arch()
         self.optimizer = optax.adamw(learning_rate=self.learning_rate)
 
@@ -246,7 +245,7 @@ def init_arch(nn_init, n_inputs_x=1):
 
     def set_arch(self):
         model = FeedForwardModule(n_layers=self.n_layers, n_neurons=self.n_hidden_x,
-                              n_out=self.n_out)
+                              n_out=self.n_out*2 if self.probabilistic else self.n_out)
         return model
 
     def fit(self, inputs, targets, n_epochs=None, savepath_tr_plots=None, stats_step=None, rel_tol=None):