Fix pytype error.

mlcommons · Nov 29, 2024 · 4bcc63a · 4bcc63a
1 parent 7ca7552
commit 4bcc63a
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/mlcroissant/mlcroissant/_src/datasets.py b/python/mlcroissant/mlcroissant/_src/datasets.py
@@ -223,7 +223,11 @@ def _filter_interesting_operations(self, filters: Filters | None) -> Operations:
         join_uuid = field.references.uuid
         graph = field.ctx.graph
         if join_uuid:
-            join_node = next(node for node in graph if node.uuid == join_uuid)
+            join_node = next(
+                node
+                for node in graph
+                if isinstance(node, Field) and node.uuid == join_uuid
+            )
             unneeded_nodes = [
                 node
                 for node in graph