small fixes

Author: robin-janssen

codes/surrogates/LatentPolynomial/latent_poly.py

@@ -399,7 +399,7 @@ def total_loss(
         traj_loss = criterion(x_pred, x_true)
 
         # identity loss (reconstruct x0)
-        identity = self.identity_loss(x_true, params)
+        identity = self.identity_loss(x_true, params.to(self.device))
 
         # derivative losses: compute once
         d_pred = self.first_derivative(x_pred)

codes/utils/data_utils.py

@@ -159,6 +159,9 @@ def check_and_load_data(
             elif normalisation_mode == "standardise":
                 data_info["mean_params"] = params_info["mean"]
                 data_info["std_params"] = params_info["std"]
+            if verbose:
+                print("Parameters normalized.")
+                print(f"Parameters info: {params_info}")
 
         data_info["log10_transform"] = True if log else False
 

run_tuning.py

@@ -5,7 +5,6 @@
 from pathlib import Path
 
 import optuna
-from optuna.study import MaxTrialsCallback
 from optuna.trial import TrialState
 from tqdm import tqdm