-- added testcase for saving NODE

manujosephv · Mar 18, 2021 · 82a30fe · 82a30fe
1 parent 6b892b6
commit 82a30fe
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/tests/test_common.py b/tests/test_common.py
@@ -15,9 +15,10 @@
 )
 
 MODEL_CONFIG_SAVE_TEST = [
-    CategoryEmbeddingModelConfig,
-    AutoIntConfig,
-    TabNetModelConfig,
+    (CategoryEmbeddingModelConfig, dict(layers="10-20")),
+    (AutoIntConfig, dict(num_heads=1,num_attn_blocks=1,)),
+    (NodeConfig, dict(num_trees=100, depth=2)),
+    (TabNetModelConfig, dict(n_a=2, n_d=2)),
 ]
 
 MODEL_CONFIG_FEATURE_EXT_TEST = [
@@ -67,7 +68,8 @@ def test_save_load(
         continuous_cols=continuous_cols,
         categorical_cols=categorical_cols,
     )
-    model_config_params = dict(task="regression")
+    model_config_class, model_config_params = model_config_class
+    model_config_params['task']="regression"
     model_config = model_config_class(**model_config_params)
     trainer_config = TrainerConfig(
         max_epochs=3, checkpoints=None, early_stopping=None, gpus=0, fast_dev_run=True