Update tsagcn.py

benedekrozemberczki · Sep 4, 2022 · e0ad1fa · e0ad1fa
1 parent d015a3c
commit e0ad1fa
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_geometric_temporal/nn/attention/tsagcn.py b/torch_geometric_temporal/nn/attention/tsagcn.py
@@ -27,7 +27,7 @@ def __init__(self, edge_index: list, num_nodes: int):
         self.A = self.get_spatial_graph(self.num_nodes)
 
     def get_spatial_graph(self, num_nodes):
-        self_mat = torch.eye(num_nodes).to(self.device)
+        self_mat = torch.eye(num_nodes)
         inward_mat = torch.squeeze(to_dense_adj(self.edge_index))
         inward_mat_norm = F.normalize(inward_mat, dim=0, p=1)
         outward_mat = inward_mat.transpose(0, 1)