Fix error of indexing shape in Optical Flow (#1087)

* Optical Flow error fix. Signed-off-by: Michał Szołucha <[email protected]> * Update dali/kernels/tensor_shape.h Co-Authored-By: Michał Zientkiewicz <[email protected]>
NVIDIA · Jul 18, 2019 · cacf792 · cacf792
1 parent 0196b27
commit cacf792
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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);