[tests] for STE

bigdata-ustc · Jun 27, 2021 · 139139c · 139139c
1 parent 9b48285
commit 139139c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/dina/gd/test_gddina.py b/tests/dina/gd/test_gddina.py
@@ -1,12 +1,14 @@
 # coding: utf-8
 # 2021/4/23 @ tongshiwei
 
+import pytest
 from EduCDM import GDDINA
 
 
-def test_train(data, conf, tmp_path):
+@pytest.mark.parametrize("ste", [True, False])
+def test_train(data, conf, tmp_path, ste):
     user_num, item_num, knowledge_num = conf
-    cdm = GDDINA(user_num, item_num, knowledge_num)
+    cdm = GDDINA(user_num, item_num, knowledge_num, ste=ste)
     cdm.train(data, test_data=data, epoch=2)
     filepath = tmp_path / "dina.params"
     cdm.save(filepath)