fix tests

mindee · Apr 27, 2022 · d42214e · d42214e
1 parent 7877ac5
commit d42214e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/pytorch/test_datasets_pt.py b/tests/pytorch/test_datasets_pt.py
@@ -24,11 +24,11 @@ def _validate_dataset(ds, input_size, batch_size=2, class_indices=False, is_poly
     else:
         assert target['boxes'].ndim == 2 and target['boxes'].shape[1:] == (4,)
     assert np.all(np.logical_and(target['boxes'] <= 1, target['boxes'] >= 0))
-    assert len(target['labels']) == len(target['boxes'])
     if class_indices:
         assert isinstance(target['labels'], np.ndarray) and target['labels'].dtype == np.int64
     else:
         assert isinstance(target['labels'], list) and all(isinstance(s, str) for s in target['labels'])
+    assert len(target['labels']) == len(target['boxes'])
 
     # Check batching
     loader = DataLoader(

diff --git a/tests/tensorflow/test_datasets_tf.py b/tests/tensorflow/test_datasets_tf.py
@@ -24,11 +24,11 @@ def _validate_dataset(ds, input_size, batch_size=2, class_indices=False, is_poly
     else:
         assert target['boxes'].ndim == 2 and target['boxes'].shape[1:] == (4,)
     assert np.all(np.logical_and(target['boxes'] <= 1, target['boxes'] >= 0))
-    assert len(target['labels']) == len(target['boxes'])
     if class_indices:
         assert isinstance(target['labels'], np.ndarray) and target['labels'].dtype == np.int64
     else:
         assert isinstance(target['labels'], list) and all(isinstance(s, str) for s in target['labels'])
+    assert len(target['labels']) == len(target['boxes'])
 
     # Check batching
     loader = DataLoader(ds, batch_size=batch_size)