Update MCD.py

WebGao · Jul 8, 2021 · 805fabe · 805fabe
1 parent 36bc1ea
commit 805fabe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/EduCDM/MCD/MCD.py b/EduCDM/MCD/MCD.py
@@ -59,7 +59,7 @@ def train(self, train_data, test_data=None, *, epoch: int, device="cpu", lr=0.00
             print("[Epoch %d] LogisticLoss: %.6f" % (e, float(np.mean(losses))))
 
             if test_data is not None:
-                auc, accuracy = self.eval(test_data)
+                auc, accuracy = self.eval(test_data, device=device)
                 print("[Epoch %d] auc: %.6f, accuracy: %.6f" % (e, auc, accuracy))
 
     def eval(self, test_data, device="cpu") -> tuple: