[Minor] set as validation one by default

ai4co · Sep 11, 2023 · 90a2246 · 90a2246
1 parent 4ed3235
commit 90a2246
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rl4co/models/rl/common/base.py b/rl4co/models/rl/common/base.py
@@ -133,7 +133,7 @@ def setup(self, stage="fit"):
         )
         self.train_batch_size = train_bs
         self.val_batch_size = train_bs if val_bs is None else val_bs
-        self.test_batch_size = train_bs if test_bs is None else test_bs
+        self.test_batch_size = self.val_batch_size if test_bs is None else test_bs
 
         log.info("Setting up datasets")