From c0a0d0a73bab1aaa30d7c27b4997aba5b1b57430 Mon Sep 17 00:00:00 2001 From: rohitgr7 Date: Tue, 12 Apr 2022 14:17:09 +0530 Subject: [PATCH] Update deepspeed precision test --- tests/strategies/test_deepspeed_strategy.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/strategies/test_deepspeed_strategy.py b/tests/strategies/test_deepspeed_strategy.py index c99cd2fe357bd..a9b30b0015bf3 100644 --- a/tests/strategies/test_deepspeed_strategy.py +++ b/tests/strategies/test_deepspeed_strategy.py @@ -168,12 +168,13 @@ def test_deepspeed_strategy_env(tmpdir, monkeypatch, deepspeed_config): @RunIf(deepspeed=True) +@mock.patch("torch.cuda.device_count", return_value=1) @pytest.mark.parametrize("precision", [16, "mixed"]) @pytest.mark.parametrize( "amp_backend", ["native", pytest.param("apex", marks=RunIf(amp_apex=True))], ) -def test_deepspeed_precision_choice(amp_backend, precision, tmpdir): +def test_deepspeed_precision_choice(_, amp_backend, precision, tmpdir): """Test to ensure precision plugin is also correctly chosen. DeepSpeed handles precision via Custom DeepSpeedPrecisionPlugin