class_docstr: PipelineHelper assists with implementing pipelining in model parallelism.
  
  For a tutorial on model parallelism, as it's implemented in parts of ParlAI,
  see https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html.
  
  Usage:
  >>> my_model = PipelineHelper().make_parallel(my_model)
  
  Note that you will need to manually implement logic which handles the
  moved layers.
class_docstr_tok: ["PipelineHelper","assists","with","implementing","pipelining","in","model","parallelism",".","For","a","tutorial","on","model","parallelism",",","as","it","'s","implemented","in","parts","of","ParlAI",",","see","https",":","\/\/pytorch.org\/tutorials\/intermediate\/model_parallel_tutorial.html",".","Usage",":",">",">",">","my_model","=","PipelineHelper","(",")",".make_parallel","(","my_model",")","Note","that","you","will","need","to","manually","implement","logic","which","handles","the","moved","layers","."]
docstr: Estimate the number of chunks we should split the batch into via heuristics.
docstr_tok: ["Estimate","the","number","of","chunks","we","should","split","the","batch","into","via","heuristics","."]
code: def guess_split_size(
      item: Chunk, num_gpus: Optional[int] = None, dim=0
  ) -> int:
      """
      Estimate the number of chunks we should split the batch into via heuristics.
      """
      if num_gpus is None:
          num_gpus = torch.cuda.device_count()  # type: ignore
  
      if isinstance(item, torch.Tensor):
          if num_gpus == 1:
              # no point in chunking if we're not really doing model parallel
              return item.size(dim)
          # heuristic: use the same number of chunks as 2 * num_gpus.  this
          # isn't perfect (it ideally would be tuned differently for every model
          # and number of GPUs), but it seems to work reasonably wellenough in several
          # architectures tested.
          return max(1, item.size(dim) // int(num_gpus * 2))
      elif isinstance(item, tuple):
          return PipelineHelper.guess_split_size(item[0], num_gpus)
      elif isinstance(item, dict):
          return PipelineHelper.guess_split_size(
              list(item.values())[0], num_gpus
          )
      raise TypeError(f"Cannot determine split size for {type(item)}")
  
imports: import os
  from typing import Union, Optional, Tuple, Any, List, Sized, TypeVar
  import itertools
  from collections import namedtuple
  import parlai.utils.logging
  import parlai.utils.io
  import torch.optim
label_desc: Use image node_id 6 for calling the PipelineHelper obj's underlying member method code with example usage: obj.guess_split_size(item, num_gpus, dim) and returns: max, item, PipelineHelper, PipelineHelper