simplify statement for pytorch file_type extension
This commit is contained in:
		| @@ -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: | ||||
|   | ||||
| @@ -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}") | ||||
|   | ||||
		Reference in New Issue
	
	Block a user