docstr: This class is used to sample tasks in a distributed setting such as DDP with multiple workers.
  
  This won't work as expected if `num_workers = 0` and several dataloaders
  are being iterated on at the same time.
  
  Args:
      taskset: Dataset used to sample task.
      global_rank: Rank of the current process.
      world_size: Total of number of processes.
      num_workers: Number of workers to be provided to the DataLoader.
      epoch_length: The expected epoch length. This requires to be divisible by (num_workers * world_size).
      seed: The seed will be used on __iter__ call and should be the same for all processes.
docstr_tok: ["This","class","is","used","to","sample","tasks","in","a","distributed","setting","such","as","DDP","with","multiple","workers",".","This","wo","n't","work","as","expected","if","`","num_workers","=","0","`","and","several","dataloaders","are","being","iterated","on","at","the","same","time",".","Args",":","taskset",":","Dataset","used","to","sample","task",".","global_rank",":","Rank","of","the","current","process",".","world_size",":","Total","of","number","of","processes",".","num_workers",":","Number","of","workers","to","be","provided","to","the","DataLoader",".","epoch_length",":","The","expected","epoch","length",".","This","requires","to","be","divisible","by","(","num_workers","*","world_size",")",".","seed",":","The","seed","will","be","used","on","__iter__","call","and","should","be","the","same","for","all","processes","."]
code: def __init__(
      self,
      taskset: Any,
      global_rank: int,
      world_size: int,
      num_workers: int,
      epoch_length: int,
      seed: int,
      requires_divisible: bool = True,
  ):
      """This class is used to sample tasks in a distributed setting such as DDP with multiple workers.
  
      This won't work as expected if `num_workers = 0` and several dataloaders
      are being iterated on at the same time.
  
      Args:
          taskset: Dataset used to sample task.
          global_rank: Rank of the current process.
          world_size: Total of number of processes.
          num_workers: Number of workers to be provided to the DataLoader.
          epoch_length: The expected epoch length. This requires to be divisible by (num_workers * world_size).
          seed: The seed will be used on __iter__ call and should be the same for all processes.
  
      """
      self.taskset = taskset
      self.global_rank = global_rank
      self.world_size = world_size
      self.num_workers = 1 if num_workers == 0 else num_workers
      self.worker_world_size = self.world_size * self.num_workers
      self.epoch_length = epoch_length
      self.seed = seed
      self.iteration = 0
      self.requires_divisible = requires_divisible
      self.counter = 0
  
      if (
          requires_divisible
          and epoch_length % self.worker_world_size != 0
      ):
          raise Exception(
              "The `epoch_length` should be divisible by `world_size`."
          )
  
file_path: lightning-flash/src/flash/image/classification/integrations/learn2learn.py
imports: from typing import Any, Callable, Optional
  import pytorch_lightning
  from torch.utils.data import IterableDataset
  from torch.utils.data._utils.worker import get_worker_info
  from flash.core.data.utilities.collate import default_collate
  from flash.core.utilities.imports import requires
label_desc: Use image node_id 1 to create a new TaskDistributedDataParallel object from inherited base classes: IterableDataset with example: obj = TaskDistributedDataParallel(taskset, global_rank, world_size, num_workers, epoch_length, seed, requires_divisible)