add hist model

bell99431 · Feb 7, 2024 · a84963c · a84963c
1 parent 9726a8d
commit a84963c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/qlib/contrib/model/pytorch_hist.py b/qlib/contrib/model/pytorch_hist.py
@@ -286,7 +286,7 @@ def fit(
 
         if self.model_path is not None:
             self.logger.info("Loading pretrained model...")
-            pretrained_model.load_state_dict(torch.load(self.model_path))
+            pretrained_model.load_state_dict(torch.load(self.model_path, map_location=torch.device('cpu')))
 
         model_dict = self.HIST_model.state_dict()
         pretrained_dict = {
@@ -429,7 +429,8 @@ def cal_cos_similarity(self, x, y):  # the 2nd dimension of x and y are the same
         return cos_similarity
 
     def forward(self, x, concept_matrix):
-        device = torch.device(torch.get_device(x))
+        # device = torch.device(torch.get_device(x))
+        device = torch.device("cpu")
 
         x_hidden = x.reshape(len(x), self.d_feat, -1)  # [N, F, T]
         x_hidden = x_hidden.permute(0, 2, 1)  # [N, T, F]

diff --git a/setup.py b/setup.py
@@ -70,7 +70,7 @@ def get_version(rel_path: str) -> str:
     "lightgbm>=3.3.0",
     "tornado",
     "joblib>=0.17.0",
-    "ruamel.yaml>=0.16.12",
+    "ruamel.yaml==0.17.21",
     "pymongo==3.7.2",  # For task management
     "scikit-learn>=0.22",
     "dill",