Add tests for pytorch ragged iterator

Signed-off-by: Szymon Karpiński <skarpinski@nvidia.com>
NVIDIA · Jan 16, 2024 · e329611 · e329611
1 parent 60c54fb
commit e329611
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/dali/test/python/checkpointing/test_dali_checkpointing_fw_iterators.py b/dali/test/python/checkpointing/test_dali_checkpointing_fw_iterators.py
@@ -199,3 +199,12 @@ def get_fw_iterator_class(self):
 
  def equal(self, a, b):
  return (a == b).all()
+
+class TestPytorchRagged(FwTestBase):
+ def get_fw_iterator_class(self):
+ from nvidia.dali.plugin.pytorch import DALIRaggedIterator
+
+ return DALIRaggedIterator
+
+ def equal(self, a, b):
+ return (a == b).all()
diff --git a/qa/TL0_python-self-test-core/test_body.sh b/qa/TL0_python-self-test-core/test_body.sh
@@ -41,6 +41,7 @@ test_type_annotations() {
 test_pytorch() {
  ${python_invoke_test} --attr '!slow,pytorch' test_dali_variable_batch_size.py
  ${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
  ${python_new_invoke_test} -A 'pytorch' -s type_annotations
 }