Apply suggestions from review

Signed-off-by: Rostan Tabet <rtabet@nvidia.com>

Author: rostan-t

dali/python/nvidia/dali/experimental/dynamic/_invocation.py

@@ -91,7 +91,7 @@ def __init__(
         self._future: Optional[_Future] = None
         self._run_lock = threading.Lock()
         if caller_depth is None:
-            caller_depth = 2 if getattr(self._operator, "_is_reader", False) else 4
+            caller_depth = 3 if getattr(self._operator, "_is_reader", False) else 4
         self._call_stack = (
             capture_stack(caller_depth + 1)
             if _EvalMode.current().value <= _EvalMode.eager.value

dali/python/nvidia/dali/experimental/dynamic/_op_builder.py

@@ -203,6 +203,8 @@ def init(self, max_batch_size, name, **kwargs):
         op_class.__base__.__init__(self, max_batch_size, name, **kwargs)
         if is_reader:
             self._tensor_args = {k: v for k, v in tensor_kwargs.items() if v is not None}
+            if any(isinstance(v, Batch) for v in self._tensor_args.values()):
+                raise ValueError("Readers cannot be constructed with batch keyword arguments")
         if stateful:
             self._call_id = 0
 
@@ -287,11 +289,8 @@ def call(self, *raw_args, batch_size=None, _process_params=True, **raw_kwargs):
             if overlap:
                 raise ValueError(
                     f"Keyword argument{'s'[:len(overlap)^1]} {sorted(overlap)}"
-                    f" cannot be passed both in the constructor and __call__."
+                    f" cannot be passed in both in the constructor and __call__."
                 )
-            for arg in self._tensor_args.values():
-                if isinstance(arg, Batch):
-                    raise ValueError("Readers cannot be constructed with batch keyword arguments")
             raw_kwargs = {**raw_kwargs, **self._tensor_args}
 
         batch_size = _ops._infer_batch_size(batch_size, *raw_args, **raw_kwargs)

dali/python/nvidia/dali/experimental/dynamic/_ops.py

@@ -668,7 +668,7 @@ def next_epoch(self, batch_size=None, ctx: _eval_context.EvalContext | None = No
         else:
             return self._samples(ctx)
 
-    def _process_tensor_args(self, batch_size):
+    def _process_tensor_args(self, batch_size: int | None):
         """Converts stored tensor args to Batch/Tensor form for the given batch_size."""
         if not self._tensor_args:
             return {}
@@ -699,8 +699,10 @@ def _samples(self, ctx: _eval_context.EvalContext | None = None):
                     self._actual_batch_size = 1
                 if self._max_batch_size is None:
                     self._max_batch_size = self._actual_batch_size
-                self._init_backend(ctx, (), self._process_tensor_args(self._actual_batch_size))
-            tensor_args = self._process_tensor_args(self._actual_batch_size)
+                tensor_args = self._process_tensor_args(self._actual_batch_size)
+                self._init_backend(ctx, (), tensor_args)
+            else:
+                tensor_args = self._process_tensor_args(self._actual_batch_size)
             meta = self._op_backend.GetReaderMeta()
             idx = 0
             padded_size = meta["epoch_size_padded"]