adelorenzo/stable
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

simplify statement for pytorch file_type extension

Browse Source
This commit is contained in:
Yinon Polak 2023-03-13 00:16:44 +02:00
parent 0012fe36ca
commit 523a58d3d6
2 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
freqtrade/freqai/data_drawer.py
View File

@ -446,8 +446,7 @@ class FreqaiDataDrawer:
dump(model, save_path / f"{dk.model_filename}_model.joblib") dump(model, save_path / f"{dk.model_filename}_model.joblib")
elif self.model_type == 'keras': elif self.model_type == 'keras':
model.save(save_path / f"{dk.model_filename}_model.h5") model.save(save_path / f"{dk.model_filename}_model.h5")
elif ('stable_baselines' in self.model_type or elif self.model_type in ["stable_baselines", "sb3_contrib", "pytorch"]:
self.model_type in ['sb3_contrib', 'pytorch']):
model.save(save_path / f"{dk.model_filename}_model.zip") model.save(save_path / f"{dk.model_filename}_model.zip")
if dk.svm_model is not None: if dk.svm_model is not None:

3
freqtrade/freqai/freqai_interface.py
View File

@ -563,8 +563,7 @@ class IFreqaiModel(ABC):
file_type = ".joblib" file_type = ".joblib"
elif self.dd.model_type == 'keras': elif self.dd.model_type == 'keras':
file_type = ".h5" file_type = ".h5"
elif ('stable_baselines' in self.dd.model_type or elif self.dd.model_type in ["stable_baselines", "sb3_contrib", "pytorch"]:
self.dd.model_type in ['sb3_contrib', 'pytorch']):
file_type = ".zip" file_type = ".zip"
path_to_modelfile = Path(dk.data_path / f"{dk.model_filename}_model{file_type}") path_to_modelfile = Path(dk.data_path / f"{dk.model_filename}_model{file_type}")