fix: torch_forecasting_model load_weights with float32

unit8co · Nov 1, 2023 · d09ac98 · d09ac98
1 parent ea37dc9
commit d09ac98
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/darts/models/forecasting/torch_forecasting_model.py b/darts/models/forecasting/torch_forecasting_model.py
@@ -1872,8 +1872,9 @@ def load_weights_from_checkpoint(
         )
 
         # pl_forecasting module saves the train_sample shape, must recreate one
+        np_dtype = np.float32 if ckpt["model_dtype"] == torch.float32 else np.float64
         mock_train_sample = [
-            np.zeros(sample_shape) if sample_shape else None
+            np.zeros(sample_shape, dtype=np_dtype) if sample_shape else None
             for sample_shape in ckpt["train_sample_shape"]
         ]
         self.train_sample = tuple(mock_train_sample)