Fixup for multiple entry point functions.

iree-org · Jul 2, 2024 · f196208 · f196208
1 parent 3e678c7
commit f196208
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/shark_turbine/transforms/general/add_metadata.py b/shark_turbine/transforms/general/add_metadata.py
@@ -50,8 +50,9 @@ def parse_metadata_dict(metadata_dict: dict) -> DictAttr:
 
         metadata_dict_attr = parse_metadata_dict(self.inp_metadata)
         for func_op in self.funcs:
-            if func_op.op.attributes[1].attr.value == self.func_name:
-                func_op.op.attributes["iree.reflection"] = metadata_dict_attr
+            for attr in func_op.op.attributes:
+                if isinstance(attr.attr, StringAttr) and attr.attr.value == self.func_name:
+                    func_op.op.attributes["iree.reflection"] = metadata_dict_attr
         return self.mlir_module