adjust metrics input

autogl/module/train/evaluation.py

@@ -120,6 +120,7 @@ class Auc(Evaluation):
         if len(predict.shape) == 1:
             pos_predict = predict
         else:
+            assert predict.shape[1] == 2, "Cannot use auc on given data with %d classes!" % (predict.shape[1])
             pos_predict = predict[:, 1]
         return roc_auc_score(label, pos_predict)
 
@@ -142,7 +143,11 @@ class Acc(Evaluation):
         """
         Should return: the evaluation result (float)
         """
-        return accuracy_score(label, np.argmax(predict, axis=1))
+        if len(predict.shape) == 2:
+            predict = np.argmax(predict, axis=1)
+        else:
+            predict = [1 if p > 0.5 else 0 for p in predict]
+        return accuracy_score(label, predict)
 
 
 @register_evaluate("mrr")
@@ -163,5 +168,9 @@ class Mrr(Evaluation):
         """
         Should return: the evaluation result (float)
         """
-        pos_predict = predict[:, 1]
+        if len(predict.shape) == 2:
+            assert predict.shape[1] == 2, "Cannot use mrr on given data with %d classes!" % (predict.shape[1])
+            pos_predict = predict[:, 1]
+        else:
+            pos_predict = predict
         return label_ranking_average_precision_score(label, pos_predict)