From 3dee2759ee795a3bcbaa60fb863de5ee03b161a1 Mon Sep 17 00:00:00 2001 From: Jv Kyle Eclarin Date: Sun, 16 Jan 2022 11:13:37 -0500 Subject: [PATCH] update tests for v2 (#11487) --- tests/loops/test_all.py | 3 ++- tests/loops/test_evaluation_loop.py | 7 ++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/loops/test_all.py b/tests/loops/test_all.py index 5975937018..f73c6748a5 100644 --- a/tests/loops/test_all.py +++ b/tests/loops/test_all.py @@ -82,7 +82,8 @@ def test_callback_batch_on_device(tmpdir): limit_val_batches=1, limit_test_batches=1, limit_predict_batches=1, - gpus=1, + accelerator="gpu", + devices=1, callbacks=[batch_callback], ) trainer.fit(model) diff --git a/tests/loops/test_evaluation_loop.py b/tests/loops/test_evaluation_loop.py index 30095beca2..e8ca0f195f 100644 --- a/tests/loops/test_evaluation_loop.py +++ b/tests/loops/test_evaluation_loop.py @@ -100,7 +100,12 @@ def test_memory_consumption_validation(tmpdir): torch.cuda.empty_cache() trainer = Trainer( - gpus=1, default_root_dir=tmpdir, fast_dev_run=2, move_metrics_to_cpu=True, enable_model_summary=False + accelerator="gpu", + devices=1, + default_root_dir=tmpdir, + fast_dev_run=2, + move_metrics_to_cpu=True, + enable_model_summary=False, ) trainer.fit(BoringLargeBatchModel())