Hotfix DatasetEntity.get_combined_subset function loop (#2577)

Fix get_combined_subset function
openvinotoolkit · Oct 25, 2023 · 1285168 · 1285168
1 parent 0981035
commit 1285168
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/otx/api/entities/datasets.py b/src/otx/api/entities/datasets.py
@@ -349,8 +349,9 @@ def get_combined_subset(self, subsets: List[Subset]) -> "DatasetEntity":
         Returns:
             DatasetEntity: DatasetEntity with items matching subsets
         """
+        to_keep = set(subsets)
         dataset = DatasetEntity(
-            items=[item for item in self._items if item.subset in set(subsets)],
+            items=[item for item in self if item.subset in to_keep],
             purpose=self.purpose,
         )
         return dataset