Merge pull request #1217 from chenyuwuxin/master

FIX: code optimization of SRGNN
RUCAIBox · Mar 30, 2022 · 4f76169 · 4f76169
2 parents 430db16 + 2cc57ea
commit 4f76169
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/recbole/model/sequential_recommender/srgnn.py b/recbole/model/sequential_recommender/srgnn.py
@@ -177,7 +177,7 @@ def _get_slice(self, item_seq):
         # The relative coordinates of the item node, shape of [batch_size, max_session_len]
         alias_inputs = torch.LongTensor(alias_inputs).to(self.device)
         # The connecting matrix, shape of [batch_size, max_session_len, 2 * max_session_len]
-        A = torch.FloatTensor(A).to(self.device)
+        A = torch.FloatTensor(np.array(A)).to(self.device)
         # The unique item nodes, shape of [batch_size, max_session_len]
         items = torch.LongTensor(items).to(self.device)