From ba90d16a1776ccf5a3e04c08c83a7a27e241fa71 Mon Sep 17 00:00:00 2001 From: wangchunlin Date: Fri, 10 Jan 2025 08:21:42 +0000 Subject: [PATCH] add weights_only=False, deal torch version warning --- utils/model_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/model_trainer.py b/utils/model_trainer.py index bb38684..6f360e1 100644 --- a/utils/model_trainer.py +++ b/utils/model_trainer.py @@ -111,7 +111,7 @@ class MLModel: else: raise FileNotFoundError(f"No model found in either path: {model_path} or {default_model_path}") - self.model.load_state_dict(torch.load(model_path, map_location=self.device)) + self.model.load_state_dict(torch.load(model_path, map_location=self.device, weights_only=False)) def plot_correlation(self, X: np.ndarray, feature_names: List[str]) -> None: """绘制特征相关性矩阵"""