corrected n_out bug with optuna

supsi-dacd-isaac · Dec 12, 2023 · c057b40 · c057b40
1 parent 2d5cc09
commit c057b40
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/pyforecaster/forecasting_models/neural_forecasters.py b/pyforecaster/forecasting_models/neural_forecasters.py
@@ -542,7 +542,8 @@ def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_pat
 
     def set_arch(self):
         model = PartiallyICNN(num_layers=self.n_layers, features_x=self.n_hidden_x, features_y=self.n_hidden_y,
-                              features_out=self.n_out, init_type=self.init_type, augment_ctrl_inputs=self.augment_ctrl_inputs, probabilistic=self.probabilistic)
+                              features_out=self.n_out*2 if self.probabilistic else self.n_out, init_type=self.init_type,
+                              augment_ctrl_inputs=self.augment_ctrl_inputs, probabilistic=self.probabilistic)
         return model
 
     @staticmethod
@@ -650,8 +651,8 @@ def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_pat
 
     def set_arch(self):
         model = PartiallyIQCNN(num_layers=self.n_layers, features_x=self.n_hidden_x, features_y=self.n_hidden_y,
-                              features_out=self.n_out, init_type=self.init_type, augment_ctrl_inputs=self.augment_ctrl_inputs,
-                               probabilistic=self.probabilistic)
+                              features_out=self.n_out*2 if self.probabilistic else self.n_out, init_type=self.init_type,
+                               augment_ctrl_inputs=self.augment_ctrl_inputs, probabilistic=self.probabilistic)
         return model
 
 
@@ -678,7 +679,7 @@ def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_pat
 
     def set_arch(self):
         model = PartiallyICNN(num_layers=self.n_layers, features_x=self.n_hidden_x, features_y=self.n_hidden_y,
-                              features_out=self.n_out, init_type=self.init_type,
+                              features_out=self.n_out*2 if self.probabilistic else self.n_out, init_type=self.init_type,
                                augment_ctrl_inputs=self.augment_ctrl_inputs, activation=nn.sigmoid,
                                rec_activation=nn.sigmoid, probabilistic=self.probabilistic)
         return model
@@ -702,6 +703,6 @@ def __init__(self, learning_rate: float = 0.01, batch_size: int = None, load_pat
 
     def set_arch(self):
         model = PartiallyICNN(num_layers=self.n_layers, features_x=self.n_hidden_x, features_y=self.n_hidden_y,
-                              features_out=self.n_out, activation=nn.relu, init_type=self.init_type,
-                              probabilistic=self.probabilistic)
+                              features_out=self.n_out*2 if self.probabilistic else self.n_out, activation=nn.relu,
+                              init_type=self.init_type, probabilistic=self.probabilistic)
         return model