Add check for tests that need Nvidia GPU and cuda

To enable tests to run non Nvidia GPU environments, need
to check for cuda

Signed-off-by: Martin Hickey <[email protected]>

Author: hickeyma

.gitignore

@@ -37,3 +37,6 @@ venv/
 # generated by setuptools_scm
 /fms_mo/_version.py
 
+#Generated by tests
+qcfg.json
+

tests/models/conftest.py

@@ -505,52 +505,54 @@ def model_config_fp16():
     return deepcopy(ToyModel4().half())
 
 
-class ToyModelQuantized(torch.nn.Module):
-    """
-    Three layer Linear model that has a quantized layer
+# QLinear class requires Nvidia GPU and cuda
+if torch.cuda.is_available():
+    class ToyModelQuantized(torch.nn.Module):
+        """
+        Three layer Linear model that has a quantized layer
 
-    Extends:
-        torch.nn.Module
-    """
+        Extends:
+            torch.nn.Module
+        """
 
-    def __init__(self):
-        super().__init__()
-        kwargs = {"qcfg": qconfig_init()}  # QLinear requires qconfig to work
-        self.first_layer = torch.nn.Linear(3, 3, bias=True)
-        self.second_layer = QLinear(3, 3, bias=True, **kwargs)
-        self.third_layer = torch.nn.Linear(3, 3, bias=True)
+        def __init__(self):
+            super().__init__()
+            kwargs = {"qcfg": qconfig_init()}  # QLinear requires qconfig to work
+            self.first_layer = torch.nn.Linear(3, 3, bias=True)
+            self.second_layer = QLinear(3, 3, bias=True, **kwargs)
+            self.third_layer = torch.nn.Linear(3, 3, bias=True)
 
-    def forward(self, input_tensor):
-        """
-        Forward func for Toy Model
+        def forward(self, input_tensor):
+            """
+            Forward func for Toy Model
 
-        Args:
-            input_tensor (torch.FloatTensor): Tensor to operate on
+            Args:
+                input_tensor (torch.FloatTensor): Tensor to operate on
 
-        Returns:
-            torch.FloatTensor:
-        """
-        out = self.first_layer(input_tensor)
-        out = self.second_layer(out)
-        out = self.third_layer(out)
-        return out
+            Returns:
+                torch.FloatTensor:
+            """
+            out = self.first_layer(input_tensor)
+            out = self.second_layer(out)
+            out = self.third_layer(out)
+            return out
 
 
-model_quantized_params = [ToyModelQuantized()]
+    model_quantized_params = [ToyModelQuantized()]
 
 
-@pytest.fixture(scope="function", params=model_quantized_params)
-def model_quantized(request):
-    """
-    Toy Model that has quantized layer
+    @pytest.fixture(scope="function", params=model_quantized_params)
+    def model_quantized(request):
+        """
+        Toy Model that has quantized layer
 
-    Args:
-        request (torch.nn.Module): Toy Model
+        Args:
+            request (torch.nn.Module): Toy Model
 
-    Returns:
-        torch.nn.Module: Toy Model
-    """
-    return deepcopy(request.param)
+        Returns:
+            torch.nn.Module: Toy Model
+        """
+        return deepcopy(request.param)
 
 
 # Get a model to test layer uniqueness

tests/models/test_qmodelprep.py

@@ -32,22 +32,25 @@
 ################
 # Qmodel tests #
 ################
-def test_model_quantized(
-    model_quantized: torch.nn.Module,
-    sample_input_fp32: torch.FloatTensor,
-    config_fp32: dict,
-):
-    """
-    qmodel_prep should always throw RuntimeError if a model is already quantized
 
-    Args:
-        model_quantized (torch.nn.Module): Quantized Toy Model
-        sample_input_fp32 (torch.FloatTensor): Sample fp32 input for calibration.
-        config_fp32 (dict): Config w/ fp32 settings
-    """
-    delete_config()
-    with pytest.raises(RuntimeError):
-        qmodel_prep(model_quantized, sample_input_fp32, config_fp32)
+# Requires Nvidia GPU to run
+if torch.cuda.is_available():
+    def test_model_quantized(
+        model_quantized: torch.nn.Module,
+        sample_input_fp32: torch.FloatTensor,
+        config_fp32: dict,
+    ):
+        """
+        qmodel_prep should always throw RuntimeError if a model is already quantized
+
+        Args:
+            model_quantized (torch.nn.Module): Quantized Toy Model
+            sample_input_fp32 (torch.FloatTensor): Sample fp32 input for calibration.
+            config_fp32 (dict): Config w/ fp32 settings
+        """
+        delete_config()
+        with pytest.raises(RuntimeError):
+            qmodel_prep(model_quantized, sample_input_fp32, config_fp32)
 
 
 def test_double_qmodel_prep_assert(

tox.ini

@@ -4,9 +4,6 @@ minversion = 4.4
 
 [testenv]
 description = run tests (unit)
-# Use PyTorch CPU build instead of CUDA build in test envs
-setenv =
-    PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cpu
 extras =
     dev
 package = wheel