Fixes

Signed-off-by: Joaquin Anton Guirao <[email protected]>

Author: jantonguirao

dali/test/python/test_dali_variable_batch_size.py

@@ -24,7 +24,6 @@
 from nose_utils import SkipTest, attr, nottest
 from nvidia.dali.pipeline import Pipeline, pipeline_def
 from nvidia.dali.pipeline.experimental import pipeline_def as experimental_pipeline_def
-from nvidia.dali.plugin.numba.fn.experimental import numba_function
 
 import test_utils
 from segmentation_test_utils import make_batch_select_masks
@@ -375,41 +374,30 @@ def numba_setup_out_shape(out_shape, in_shape):
     (fn.ones_like, {"devices": ["cpu"]}),
 ]
 
+numba_compatible_devices = []
 if check_numba_compatibility_gpu(False):
-    ops_image_custom_args.append(
-        (
-            numba_function,
-            {
-                "batch_processing": True,
-                "devices": ["cpu"],
-                "in_types": [types.UINT8],
-                "ins_ndim": [3],
-                "out_types": [types.UINT8],
-                "outs_ndim": [3],
-                "run_fn": numba_set_all_values_to_255_batch,
-                "setup_fn": numba_setup_out_shape,
-            },
-        )
-    )
-
-
+    numba_compatible_devices.append("gpu")
 if check_numba_compatibility_cpu(False):
-    ops_image_custom_args.append(
-        (
-            numba_function,
-            {
-                "batch_processing": False,
-                "devices": ["cpu"],
-                "in_types": [types.UINT8],
-                "ins_ndim": [3],
-                "out_types": [types.UINT8],
-                "outs_ndim": [3],
-                "run_fn": numba_set_all_values_to_255_batch,
-                "setup_fn": numba_setup_out_shape,
-            },
+    numba_compatible_devices.append("cpu")
+
+if len(numba_compatible_devices) > 0:
+    from nvidia.dali.plugin.numba.fn.experimental import numba_function
+    for device in numba_compatible_devices:
+        ops_image_custom_args.append(
+            (
+                numba_function,
+                {
+                    "batch_processing": True if device == 'gpu' else False,
+                    "devices": [device],
+                    "in_types": [types.UINT8],
+                    "ins_ndim": [3],
+                    "out_types": [types.UINT8],
+                    "outs_ndim": [3],
+                    "run_fn": numba_set_all_values_to_255_batch,
+                    "setup_fn": numba_setup_out_shape,
+                },
+            )
         )
-    )
-
 
 def test_ops_image_custom_args():
     for op, args in ops_image_custom_args:

qa/TL0_self_test_Ampere/test.sh

@@ -1,6 +1,6 @@
 #!/bin/bash -ex
 
-pip_packages='${python_test_runner_package} numpy opencv-python nvidia-ml-py==11.450.51'
+pip_packages='${python_test_runner_package} numpy opencv-python nvidia-ml-py==11.450.51 numba'
 
 target_dir=./dali/test/python