🚀 Feature

Add an option to run GPU tests only, and skip all CPU tests

For example, we can have an environment variable PYTORCH_TEST_GPU_ONLY=1.

Motivation

Usually, when people only want to test GPU functionalities, they can skip CPU tests to save plenty of times.

Pitch

For example, some pytorch tests look like (assuming pytorch built with GPU)

def test_something(self, device):
    # do something

At runtime, two tests test_something_cpu and test_something_cuda will be generated. The goal is to skip the test_something_cpu test and only run test_something_cuda.

For example, we can do it like this

PYTORCH_TEST_GPU_ONLY=1 python test/run_test.py

Alternatives

N/A

How to implement

If I understand the code correctly, this feature can be easily implemented like this

1. Add TEST_GPU_ONLY = os.getenv('PYTORCH_TEST_GPU_ONLY', '0') == '1' here
   

   pytorch/torch/testing/_internal/common_utils.py

   Lines 346 to 347 in 665feda

   	TEST_SKIP_FAST = os.getenv('PYTORCH_TEST_SKIP_FAST', '0') == '1'

2. Add something like

    if TEST_GPU_ONLY:
        fullname = self.id().lower()
        is_cuda_test = 'gpu' in fullname or 'cuda' in fullname or 'cuda' in torch.tensor([]).device.type
        if not is_cuda_test:
            raise unittest.SkipTest("Test is probably not a GPU test. We disabled it with PYTORCH_TEST_GPU_ONLY")

here in setup of class common_utils.TestCase

Additional context

N/A

cc @mruberry @VitalyFedyunin @ptrblck @ngimel