fix(NeuralProphet): use their save & load method

dream-faster · Mar 23, 2023 · 9e1afda · 9e1afda
1 parent 031baac
commit 9e1afda
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/fold_models/neuralprophet.py b/src/fold_models/neuralprophet.py
@@ -1,4 +1,3 @@
-from copy import copy
 from typing import Any, Optional, Union
 
 import pandas as pd
@@ -40,16 +39,13 @@ def predict(self, X: pd.DataFrame) -> Union[pd.Series, pd.DataFrame]:
         )
 
     def __deepcopy__(self, memo):
-        model = copy(self.model)
-        setattr(model, "trainer", None)
         from io import BytesIO
 
-        from torch import load, save
+        from neuralprophet.utils import load, save
 
         buff = BytesIO()
 
-        save(model, buff)
+        save(self.model, buff)
         buff.seek(0)
         model = load(buff)
-        model.restore_trainer()
         return NeuralProphetWrapper(model)