class_docstr: A collection of remote modules (of type RemoteModule) with connections showing how inputs
  to the model or outputs of individual modules are use as inputs of subsequent modules.
  The graph has a number of helper functions that add new modules to the graph and define inputs
  to these module.
docstr: Compiles a list of modules (starting from module number module_idx), where each module in the list
  gets the output of previous module in the list as its input. So every module in the list, except the
  first one should have only one input, and similarly, every module in the list, except the last one
  should have only one output.
docstr_tok: ["Compiles","a","list","of","modules","(","starting","from","module","number","module_idx",")",",","where","each","module","in","the","list","gets","the","output","of","previous","module","in","the","list","as","its","input",".","So","every","module","in","the","list",",","except","the","first","one","should","have","only","one","input",",","and","similarly",",","every","module","in","the","list",",","except","the","last","one","should","have","only","one","output","."]
code: def _trace_modules(self, node: Node) -> List[Node]:
      """Compiles a list of modules (starting from module number module_idx), where each module in the list
      gets the output of previous module in the list as its input. So every module in the list, except the
      first one should have only one input, and similarly, every module in the list, except the last one
      should have only one output.
      """
      partition = []
      current_node = node
      while True:
          partition.append(current_node)
          # If we reached a module with multiple outputs or with multiple consumers for its output,
          # stop adding more modules to the partition.
          if len(current_node.output_consumers) != 1:
              break
          if current_node.num_outputs is not None:
              break
          # Next module to add is the only consumer of the ouput of the current module
          next_node = current_node.output_consumers[0].consumer
          # If the next module has multiple inputs, do not add it to the current partition and stop.
          if next_node.inputs != [DataSource(current_node, 0)]:
              break
          # If the next module is on a different deivce or worker, stop
          if next_node.module.on != current_node.module.on:
              break
          if next_node.module.device != current_node.module.device:
              break
          current_node = next_node
  
      return partition
  
imports: from dataclasses import dataclass
  from typing import List, Optional, Set, Tuple, Union
  from torch import Tensor, nn
  from torch.distributed import rpc
  from torch.distributed.nn import RemoteModule
  from .data import DataConsumer
label_desc: Use image node_id 8 for calling the PipelineModulesGraph obj's underlying member method code with example usage: obj._trace_modules(node) and returns: partition