Fix for OF test (NVIDIA#1008)

* fix optical flow test error Signed-off-by: Michał Szołucha <[email protected]>
phecy · Jun 25, 2019 · e09512d · e09512d
1 parent 4b1793c
commit e09512d
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/dali/pipeline/operators/optical_flow/optical_flow.h b/dali/pipeline/operators/optical_flow/optical_flow.h
@@ -124,16 +124,18 @@ class OpticalFlow : public Operator<Backend> {
       sequence_sizes_[i] = shape[i][0];
     }
 
-    for (const auto &seq : shape) {
-      DALI_ENFORCE(seq[1] != frames_height_ || seq[2] != frames_width_ || seq[3] != depth_,
-        "Width, height and depth for Optical Flow calculation must be equal for all sequences.");
+    for (auto sz : sequence_sizes_) {
+      DALI_ENFORCE(sz >= 2, (sz == 1
+                             ? "One-frame sequence encountered. Make sure that all input sequences "
+                               "for Optical Flow have at least 2 frames."
+                             : "Empty sequence encountered. Make sure that all input sequences"
+                               " for Optical Flow have at least 2 frames."));
     }
 
-    for (auto sz : sequence_sizes_) {
-      DALI_ENFORCE(sz >= 2, (sz == 1 ? "One-frame sequence encountered. "
-        "Make sure that all input sequences for Optical Flow have at least 2 frames." :
-        "Empty sequence encountered. "
-        "Make sure that all input sequences for Optical Flow have at least 2 frames."));
+    for (const auto &seq : shape) {
+      DALI_ENFORCE(seq[1] == frames_height_ && seq[2] == frames_width_ && seq[3] == depth_,
+                   "Width, height and depth for Optical Flow calculation "
+                   "must be equal for all sequences.");
     }
   }