fix calls and remove unused imports for check_device_drivers

benchmarks/tests/test_benchmark.py

@@ -129,12 +129,12 @@ pytest_benchmark_param = pytest.mark.parametrize(
         pytest.param(True, "cpu", marks=pytest.mark.skip),
         pytest.param(
             False,
-            "gpu",
+            "cuda",
             marks=pytest.mark.skipif(
-                check_device_drivers("gpu"), reason="nvidia-smi not found"
+                check_device_drivers("cuda"), reason="nvidia-smi not found"
             ),
         ),
-        pytest.param(True, "gpu", marks=pytest.mark.skip),
+        pytest.param(True, "cuda", marks=pytest.mark.skip),
         pytest.param(
             False,
             "vulkan",

tank/examples/deberta-base_tf/deberta-base_tf_test.py

@@ -1,4 +1,3 @@
-from shark.iree_utils._common import check_device_drivers, device_driver_info
 from shark.shark_inference import SharkInference
 from shark.shark_downloader import download_model
 from shark.parser import shark_args

tank/examples/opt/opt_causallm.py

@@ -3,10 +3,6 @@ import os
 import torch
 import numpy as np
 from shark_opt_wrapper import OPTForCausalLMModel
-from shark.iree_utils._common import (
-    check_device_drivers,
-    device_driver_info,
-)
 from shark.shark_inference import SharkInference
 from shark.shark_importer import import_with_fx
 from transformers import AutoTokenizer, OPTForCausalLM

tank/examples/rembert_tf/rembert_tf_test.py

@@ -1,4 +1,3 @@
-from shark.iree_utils._common import check_device_drivers, device_driver_info
 from shark.shark_inference import SharkInference
 from shark.shark_downloader import download_model
 from tank.test_utils import get_valid_test_params, shark_test_name_func

tank/examples/tapas-base_tf/tapas-base_tf_test.py

@@ -44,7 +44,7 @@ class TapasBaseModuleTest(unittest.TestCase):
         self.module_tester.create_and_check_module(dynamic, device)
 
     @pytest.mark.skipif(
-        check_device_drivers("cuda"), reason=device_driver_info("gpu")
+        check_device_drivers("cuda"), reason=device_driver_info("cuda")
     )
     def test_module_static_cuda(self):
         dynamic = False