From 523a58d3d67ee058df60de0c059053bfe307f199 Mon Sep 17 00:00:00 2001 From: Yinon Polak Date: Mon, 13 Mar 2023 00:16:44 +0200 Subject: [PATCH] simplify statement for pytorch file_type extension --- freqtrade/freqai/data_drawer.py | 3 +-- freqtrade/freqai/freqai_interface.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/freqtrade/freqai/data_drawer.py b/freqtrade/freqai/data_drawer.py index fd839ad2f..aaf9a869c 100644 --- a/freqtrade/freqai/data_drawer.py +++ b/freqtrade/freqai/data_drawer.py @@ -446,8 +446,7 @@ class FreqaiDataDrawer: dump(model, save_path / f"{dk.model_filename}_model.joblib") elif self.model_type == 'keras': model.save(save_path / f"{dk.model_filename}_model.h5") - elif ('stable_baselines' in self.model_type or - self.model_type in ['sb3_contrib', 'pytorch']): + elif self.model_type in ["stable_baselines", "sb3_contrib", "pytorch"]: model.save(save_path / f"{dk.model_filename}_model.zip") if dk.svm_model is not None: diff --git a/freqtrade/freqai/freqai_interface.py b/freqtrade/freqai/freqai_interface.py index 3d06745f9..7c45d4642 100644 --- a/freqtrade/freqai/freqai_interface.py +++ b/freqtrade/freqai/freqai_interface.py @@ -563,8 +563,7 @@ class IFreqaiModel(ABC): file_type = ".joblib" elif self.dd.model_type == 'keras': file_type = ".h5" - elif ('stable_baselines' in self.dd.model_type or - self.dd.model_type in ['sb3_contrib', 'pytorch']): + elif self.dd.model_type in ["stable_baselines", "sb3_contrib", "pytorch"]: file_type = ".zip" path_to_modelfile = Path(dk.data_path / f"{dk.model_filename}_model{file_type}")