Convert torch._C._set_grad_enabled and torch._C.set_grad_enabled to torch.set_grad_enabled (#341)

Author: xilzy

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -183,7 +183,28 @@ def _impl_unstable_to_stable_softplus(self, gm):
 
     # replace this line with modification code for task 119 (torch._C._nn.one_hot)
 
-    # replace this line with modification code for task 121 (torch._C._set_grad_enabled)
+    def _impl_unstable_to_stable_set_grad_enabled(self, gm):
+        """
+        Convert torch._C._set_grad_enabled and torch._C.set_grad_enabled to torch.set_grad_enabled
+        """
+
+        def replace_in_graph(graph_mod):
+            for node in graph_mod.graph.nodes:
+                if node.op == "call_function":
+                    if "set_grad_enabled" in str(node.target):
+                        node.target = torch.set_grad_enabled
+            graph_mod.recompile()
+
+        modules = [gm]
+        modules += [
+            m
+            for _, m in gm.named_modules()
+            if isinstance(m, torch.fx.GraphModule) and m is not gm
+        ]
+        for m in modules:
+            replace_in_graph(m)
+
+        return gm
 
     # replace this line with modification code for task 122 (torch._C._log_api_usage_once)
 

graph_net/torch/fx_graph_serialize_util.py

@@ -143,7 +143,8 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
         # replace this line with modification code for task 117 (torch._C._linalg.linalg_norm)
         (r"torch\._C\._nn\.softplus\(", "torch.nn.functional.softplus("),
         # replace this line with modification code for task 119 (torch._C._nn.one_hot)
-        # replace this line with modification code for task 121 (torch._C._set_grad_enabled)
+        (r"torch\._C\._set_grad_enabled\(", "torch.set_grad_enabled("),
+        (r"torch\._C\.set_grad_enabled\(", "torch.set_grad_enabled("),
         # replace this line with modification code for task 122 (torch._C._log_api_usage_once)
         # replace this line with modification code for task 123 (torch._C._nn.pad)
         # replace this line with modification code for task 125 (torch._C._nn.gelu)