Merge pull request #118 from justinchuby/patch-1

Fix ExportedProgram display when tensor is None
google-ai-edge · Aug 7, 2024 · 4ee14be · 4ee14be
2 parents 55cd021 + 350a579
commit 4ee14be
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/src/server/package/src/model_explorer/pytorch_exported_program_adater_impl.py b/src/server/package/src/model_explorer/pytorch_exported_program_adater_impl.py
@@ -178,14 +178,15 @@ def add_node_attrs(self, fx_node: torch.fx.node.Node, node: GraphNode):
       tensor_spec = self.inputs_map.get(fx_node.name)
       if tensor_spec:
         node.attrs.append(KeyValue(key='target', value=str(tensor_spec[0])))
-        node.attrs.append(
-            KeyValue(
-                key='__value',
-                value=self.print_tensor(
-                    tensor_spec[1], self.settings['const_element_count_limit']
-                ),
-            )
-        )
+        if (tensor := tensor_spec[1]) is not None:
+          node.attrs.append(
+              KeyValue(
+                  key='__value',
+                  value=self.print_tensor(
+                      tensor, self.settings['const_element_count_limit']
+                  ),
+              )
+          )
 
   def add_outputs_metadata(self, fx_node: torch.fx.node.Node, node: GraphNode):
     out_vals = fx_node.meta.get('val')