simplify file_type check comparisons

This commit is contained in:
Yinon Polak 2023-03-12 14:50:08 +02:00
parent f9fdf1c31b
commit cb17b36981
2 changed files with 2 additions and 4 deletions

View File

@ -447,8 +447,7 @@ class FreqaiDataDrawer:
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 ('stable_baselines' in self.model_type or
'sb3_contrib' == self.model_type or self.model_type in ['sb3_contrib', 'pytorch']):
'pytorch' == self.model_type):
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:

View File

@ -564,8 +564,7 @@ class IFreqaiModel(ABC):
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 ('stable_baselines' in self.dd.model_type or
'sb3_contrib' == self.dd.model_type or self.dd.model_type in ['sb3_contrib', 'pytorch']):
'pytorch' == self.dd.model_type):
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}")