From 6941852cd103add949142bbbaa5bb3b5b9feba3b Mon Sep 17 00:00:00 2001 From: Jirka Borovec Date: Thu, 6 Aug 2020 21:16:56 +0200 Subject: [PATCH] Apply suggestions from code review --- tests/trainer/test_dataloaders.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/trainer/test_dataloaders.py b/tests/trainer/test_dataloaders.py index c251983c6aac1e..d9e2500707fc85 100644 --- a/tests/trainer/test_dataloaders.py +++ b/tests/trainer/test_dataloaders.py @@ -277,11 +277,11 @@ def test_inf_dataloaders_with_limit_percent_batches(tmpdir, limit_train_batches, results = trainer.fit(model) assert results == 1 - assert trainer.num_training_batches == 0 if limit_train_batches == 0.0 else float('inf') - assert trainer.num_val_batches[0] == 0 if limit_val_batches == 0.0 else float('inf') + assert trainer.num_training_batches == (0 if limit_train_batches == 0.0 else float('inf')) + assert trainer.num_val_batches[0] == (0 if limit_val_batches == 0.0 else float('inf')) trainer.test(ckpt_path=None) - assert trainer.num_test_batches[0] == 0 if limit_test_batches == 0.0 else float('inf') + assert trainer.num_test_batches[0] == (0 if limit_test_batches == 0.0 else float('inf')) @pytest.mark.parametrize(['limit_train_batches', 'limit_val_batches', 'limit_test_batches'], [