Enable the tests.

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
NVIDIA · Oct 21, 2024 · 672f6f9 · 672f6f9
1 parent 8c80a22
commit 672f6f9
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/dali/test/python/dlpack/test_torch.py b/dali/test/python/dlpack/test_torch.py
@@ -2,7 +2,7 @@
 import nvidia.dali.fn as fn
 import torch
 import numpy as np
-import time
+from nose.plugins.attrib import attr
 
 
 @dali.pipeline_def(batch_size=4, num_threads=1, device_id=0, prefetch_queue_depth=2)
@@ -14,14 +14,14 @@ def _test_pipe():
     return ext + row + col
 
 
+@attr("pytorch")
 def test_dlpack():
     print("Testing dlpack")
     # get a DALI pipeline that produces batches of very large tensors
     pipe = _test_pipe(experimental_exec_dynamic=True)
     pipe.build()
     pipe.run()
 
-    results = []
     s = torch.cuda.Stream(0)
     with torch.cuda.stream(s):
         sum = torch.zeros((), dtype=torch.float32).cuda()

diff --git a/dali/test/python/dlpack/test_torch_perf.py b/dali/test/python/dlpack/test_torch_perf.py
@@ -3,6 +3,7 @@
 import torch
 import numpy as np
 import time
+from nose.plugins.attrib import attr
 
 
 @dali.pipeline_def(batch_size=32, num_threads=8, device_id=0)
@@ -80,6 +81,7 @@ def to_torch(dali_tensor):
     return (end - start) * 1e-6
 
 
+@attr("pytorch")
 def test_perf():
     """Test that DLPack zero-copy output is faster than copying."""
     dlpack_times = []

diff --git a/qa/TL0_python-self-test-core/test_body.sh b/qa/TL0_python-self-test-core/test_body.sh
@@ -52,6 +52,7 @@ test_pytorch() {
     ${python_invoke_test} --attr '!slow,pytorch' test_dali_variable_batch_size.py
     if [ -z "$DALI_ENABLE_SANITIZERS" ]; then
         ${python_new_invoke_test} -A 'pytorch' -s type_annotations
+        ${python_new_invoke_test} -A 'pytorch' -s dlpack
         ${python_new_invoke_test} -A '!slow' checkpointing.test_dali_checkpointing_fw_iterators.TestPytorch
         ${python_new_invoke_test} -A '!slow' checkpointing.test_dali_checkpointing_fw_iterators.TestPytorchRagged
     fi