fix none vars in containers (apache#32)

MasterJH5574 · Mar 4, 2024 · 20ee98d · 20ee98d
2 parents 7a4e62e + 79def73
commit 20ee98d
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/frontend/guard_tracker.py b/frontend/guard_tracker.py
@@ -249,6 +249,8 @@ def record_function(self,
                 target_state = False
             elif func.__name__ == '__exit__':
                 target_state = func.__self__.prev
+            elif func.__name__ == 'clone':
+                target_state = False
             else:
                 raise ValueError(func)
             args = [
@@ -1380,6 +1382,11 @@ def make_sub_var(value: Any, fx_node: torch.fx.Node) -> None:
                             partial.extract_code_at_start)
                         self.state.objects.add(new_scalar_var, value)
                         dyn.mark_dynamic(value, dyn.ScalarWithUnknownValue())
+                    elif value is None:
+                        new_none_var = vs.NoneVar.from_value(
+                            None, False, self.state.objects.helper_functions,
+                            None, partial.extract_code_at_start)
+                        self.state.objects.add(new_none_var, None)
                     elif isinstance(value, (tuple, list)):
                         for i, sub_value in enumerate(value):
                             sub_node = self.state.fx_graph.create_node(
@@ -1726,7 +1733,7 @@ def set_if_inplace_return() -> None:
                 return
             if hasattr(func,
                        "__name__") and func.__name__ in ("flatten_parameters",
-                                                         "numel"):
+                                                         "numel", "children"):
                 return
             print("record function in graph", func)
             self.state.record_function(