diff --git a/torch/fx/passes/split_utils.py b/torch/fx/passes/split_utils.py
index 0236d92..f3d8608 100644
--- a/torch/fx/passes/split_utils.py
+++ b/torch/fx/passes/split_utils.py
@@ -3,7 +3,7 @@
 
 import torch.fx
 from torch.fx.graph import map_arg
-from .tools_common import NodeList, NodeSet
+from .tools_common import NodeList
 from torch.fx._compatibility import compatibility
 from torch.fx.passes.utils import lift_subgraph_as_module, HolderModule
 
@@ -130,7 +130,7 @@
     all_components: List[Component] = []
 
     # Stores nodes that will be used in main graph.
-    used_in_main: NodeSet = set()
+    used_in_main: Dict[torch.fx.Node, None] = {}
 
     # Main graph after split.
     main_g = torch.fx.Graph()
@@ -208,7 +208,7 @@
                 comp.input_placeholders.append(
                     comp.graph.placeholder(x.name, type_expr=x.type)
                 )
-                used_in_main.add(x)
+                used_in_main[x] = None
 
             return comp.input_placeholders[
                 next(i for i, y in enumerate(comp.orig_inputs) if x is y)
@@ -231,7 +231,7 @@
         else:
             # All component results consumed by the output node should be
             # marked as "used in main".
-            used_in_main.add(x)
+            used_in_main[x] = None
 
     # If a node is used in main graph then we mark it as an output in the component
     # it belongs to.