From a6231241b8ddf2fd6f38f8931f42778d84e78603 Mon Sep 17 00:00:00 2001 From: Federico Berto Date: Thu, 20 Jun 2024 16:10:13 +0900 Subject: [PATCH] [Tests] less ants for faster testing --- tests/test_training.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_training.py b/tests/test_training.py index 0730923d..07d23a25 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -203,7 +203,13 @@ def test_nargnn(): @pytest.mark.skipfif("numba" not in sys.modules, reason="Numba not installed") def test_deepaco(): env = TSPEnv(generator_params=dict(num_loc=20)) - model = DeepACO(env, train_data_size=10, val_data_size=10, test_data_size=10) + model = DeepACO( + env, + train_data_size=10, + val_data_size=10, + test_data_size=10, + policy_kwargs={"n_ants": 5}, + ) trainer = RL4COTrainer( max_epochs=1, gradient_clip_val=1, devices=1, accelerator=accelerator )