diff --git a/tests/unit/common.py b/tests/unit/common.py index f26ba95de4e0..f2b2dc387d38 100644 --- a/tests/unit/common.py +++ b/tests/unit/common.py @@ -108,6 +108,10 @@ def _get_fixture_kwargs(self, request, func): return fixture_kwargs def _launch_procs(self, num_procs): + if torch.cuda.is_available() and torch.cuda.device_count() < num_procs: + pytest.skip( + f"Skipping test because not enough GPUs are available: {num_procs} required, {torch.cuda.device_count()} available" + ) mp.set_start_method('forkserver', force=True) skip_msg = mp.Queue() # Allows forked processes to share pytest.skip reason processes = []