fixed numpy based size filtering
Summary:
This bug got introduced in my [commit](https://github.com/fairinternal/fairseq-py/commit/9624f9651478bcb88022decf7e1b0685b410133b) for fast numpy based size filtering.
Pull Request resolved:
@@ -171,11 +171,11 @@ def filter_by_size(indices, dataset, max_positions, raise_exception=False):
"""
if isinstance(max_positions, float) or isinstance(max_positions, int):
if hasattr(dataset, 'sizes') and isinstance(dataset.sizes, np.ndarray):
- ignored = indices[dataset.sizes > max_positions].tolist()
- indices = indices[dataset.sizes <= max_positions]
+ ignored = indices[dataset.sizes[indices] > max_positions].tolist()
+ indices = indices[dataset.sizes[indices] <= max_positions]
elif hasattr(dataset, 'sizes') and isinstance(dataset.sizes, list) and len(dataset.sizes) == 1:
- ignored = indices[dataset.sizes[0] > max_positions].tolist()
- indices = indices[dataset.sizes[0] <= max_positions]
+ ignored = indices[dataset.sizes[0][indices] > max_positions].tolist()
+ indices = indices[dataset.sizes[0][indices] <= max_positions]
else:
indices, ignored = _filter_by_size_dynamic(indices, dataset.size, max_positions)
else:
