Fix #91: ANOMALOUS device missing

pygod-team · Nov 29, 2023 · a53d8a2 · a53d8a2
1 parent 86b63b6
commit a53d8a2
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pygod/detector/anomalous.py b/pygod/detector/anomalous.py
@@ -114,15 +114,15 @@ def decision_function(self, data, label=None):
         return self.decision_score_
 
     def process_graph(self, data):
-        x = data.x
-        s = data.s
+        x = data.x.to(self.device)
+        s = data.s.to(self.device)
 
         s = torch.max(s, s.T)
         laplacian = torch.diag(torch.sum(s, dim=1)) - s
 
-        w_init = torch.randn_like(x.T)
+        w_init = torch.randn_like(x.T).to(self.device)
         r_init = torch.inverse((1 + self.weight_decay)
-            * torch.eye(x.shape[0]) + self.gamma * laplacian) @ x
+            * torch.eye(x.shape[0]).to(self.device) + self.gamma * laplacian) @ x
 
         return x, s, laplacian, w_init, r_init