diff --git a/dali/kernels/tensor_shape.h b/dali/kernels/tensor_shape.h
index 4ee54b4602..87db2e2ddd 100644
--- a/dali/kernels/tensor_shape.h
+++ b/dali/kernels/tensor_shape.h
@@ -515,6 +515,7 @@ struct TensorListShapeBase {
   void set_tensor_shape(int64_t sample, const SampleShape &sample_shape) {
     detail::check_compatible_ndim<sample_ndim, compile_time_size<SampleShape>::value>();
     assert(static_cast<int>(dali::size(sample_shape)) == static_cast<int>(sample_dim()));
+    assert(sample >= 0 && sample < nsamples && "Sample index out of range");
     int64_t base = sample_dim() * sample;
     for (int i = 0; i < sample_dim(); i++) {
       shapes[base + i] = sample_shape[i];
diff --git a/dali/pipeline/operators/optical_flow/optical_flow.cc b/dali/pipeline/operators/optical_flow/optical_flow.cc
index 22d3b3a27b..a3ddc83de1 100644
--- a/dali/pipeline/operators/optical_flow/optical_flow.cc
+++ b/dali/pipeline/operators/optical_flow/optical_flow.cc
@@ -73,7 +73,7 @@ void OpticalFlow<GPUBackend>::RunImpl(Workspace<GPUBackend> *ws, const int) {
     kernels::TensorListShape<> new_sizes(nsequences_, 1 + out_shape.sample_dim());
     for (int i = 0; i < nsequences_; i++) {
       auto shape = kernels::shape_cat(sequence_sizes_[i] - 1, out_shape);
-      new_sizes.set_tensor_shape(1, shape);
+      new_sizes.set_tensor_shape(i, shape);
     }
     output.Resize(new_sizes);
 
@@ -114,7 +114,7 @@ void OpticalFlow<GPUBackend>::RunImpl(Workspace<GPUBackend> *ws, const int) {
     kernels::TensorListShape<> new_sizes(nsequences_, 1 + out_shape.sample_dim());
     for (int i = 0; i < nsequences_; i++) {
       auto shape = kernels::shape_cat(sequence_sizes_[i] - 1, out_shape);
-      new_sizes.set_tensor_shape(1, shape);
+      new_sizes.set_tensor_shape(i, shape);
     }
     output.Resize(new_sizes);