Back out "Back out "[const_fold] Set requires_grad based on the folded tensor; add device_for_folding option"" (#79696)

jfix71 · pytorchmergebot · commit d4bbebbb97ea · 2022-06-18T19:13:49.000Z
Summary: This is an un-backout but with a small change to set the default device `device_for_folded_attrs="cuda"` instead of `"cpu"`, which should avoid BC issues for TRT lowering. Original commit changeset: 4ae1863e28ff Original Phabricator Diff: D37192230 (24c2aff) Test Plan: Added unit test Differential Revision: D37205432 Pull Request resolved: #79696 Approved by: https://github.com/dborkovic
diff --git a/test/fx/test_fx_const_fold.py b/test/fx/test_fx_const_fold.py
@@ -693,7 +693,9 @@ def forward(self, x, y):
         gm = torch.fx.symbolic_trace(mod)
         in_x, in_y = torch.tensor([[-0.45]]), torch.tensor([0.9])
         ShapeProp(gm).propagate(in_x, in_y)
-        mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(gm)
+        mod_folded: const_fold.FoldedGraphModule = const_fold.split_const_subgraphs(
+            gm, device_for_folded_attrs="cpu"
+        )
         self._verify_const_fold_mod(mod_folded)
 
         mod_folded.run_folding()
diff --git a/torch/fx/experimental/const_fold.py b/torch/fx/experimental/const_fold.py
@@ -22,7 +22,7 @@ def __init__(
         graph: torch.fx.Graph,
         const_subgraph: Optional[torch.fx.Graph] = None,
         fx_const_folded_attrs_name: str = None,
-        device_for_folded_attrs: str = "cpu",
+        device_for_folded_attrs: str = "cuda",
     ):
         # In init, we set graph's owning module to root which will make graph's
         # owning module be None because graph already have a owning module. We
