Fix.

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>

Author: mzient

dali/python/nvidia/dali/experimental/dali2/_batch.py

@@ -18,7 +18,6 @@
     Tensor,
     _is_full_slice,
     _try_convert_enums,
-    _backend_device,
     tensor as _tensor,
     as_tensor as _as_tensor,
 )
@@ -29,6 +28,13 @@
 from . import _invocation
 import nvtx
 
+def _backend_device(backend: Union[_backend.TensorListCPU, _backend.TensorListGPU]) -> Device:
+    if isinstance(backend, _backend.TensorListCPU):
+        return Device("cpu")
+    elif isinstance(backend, _backend.TensorListGPU):
+        return Device("gpu", backend.device_id())
+    else:
+        raise ValueError(f"Unsupported backend type: {type(backend)}")
 
 def _is_tensor_type(x):
     from . import _batch