parameters: (
          self, total_size_of_graph, logical_device_id: int = 0
      )
return_statement: return
docstring: Fit the whole fx module into one device
docstring_summary: Fit the whole fx module into one device
docstring_tokens: Fit the whole fx module into one device
function: def find_single_partition(
          self, total_size_of_graph, logical_device_id: int = 0
      ) -> None:
          """Fit the whole fx module into one device"""
          partition_0 = self.create_partition()
          for node in self.graph_module.graph.nodes:
              if node.op == "output":
                  # Skip the output node, but there can
                  # be nodes after the output in certain cases.
                  continue
              partition_0.nodes.add(node)
          partition_0.used_mem_bytes = total_size_of_graph
          partition_0.logical_device_ids = [logical_device_id]
          # Get the node to partition mapping
          self.node_to_partition = get_node_to_partition_mapping(self.partitions)
          return
function_tokens: def find_single_partition ( self , total_size_of_graph , logical_device_id : int = 0 ) -> None : partition_0 = self . create_partition ( ) for node in self . graph_module . graph . nodes : if node . op == "output" : # Skip the output node, but there can # be nodes after the output in certain cases. continue partition_0 . nodes . add ( node ) partition_0 . used_mem_bytes = total_size_of_graph partition_0 . logical_device_ids = [ logical_device_id ] # Get the node to partition mapping self . node_to_partition = get_node_to_partition_mapping ( self . partitions ) return