code: def get_train_loader(
      nranks, train_dataset, batch_size, num_workers, collate_fn=None
  ):
      train_sampler = None
      is_shuffle = True
      batch_size = batch_size
  
      if nranks > 1:
          is_shuffle = False
          batch_size_per_gpu = int(batch_size / nranks)
          train_sampler = DistributedBatchSampler(
              train_dataset,
              batch_size=batch_size_per_gpu,
              shuffle=is_shuffle,
              drop_last=True,
          )
  
          train_loader = DataLoader(
              train_dataset,
              batch_sampler=train_sampler,
              num_workers=num_workers,
              use_shared_memory=True,
              collate_fn=collate_fn,
          )
      else:
          train_loader = DataLoader(
              train_dataset,
              batch_size=batch_size,
              num_workers=num_workers,
              drop_last=True,
              shuffle=is_shuffle,
              use_shared_memory=True,
              collate_fn=collate_fn,
          )
  
      return train_loader
  
file_path: PaddleSeg/contrib/CrossPseudoSupervision/core/train.py
imports: import os
  import time
  from collections import deque
  import shutil
  import paddle
  import paddle.distributed
  from paddle.io import DistributedBatchSampler, DataLoader
  from paddleseg.utils import TimeAverager, calculate_eta, logger, train_profiler, op_flops_funs
  from paddleseg.core.val import evaluate
  from utils import cps_resume
  from batch_transforms import SegCollate, AddMaskParamsToBatch
label_desc: Use image node_id 1 for calling a global function with example usage: get_train_loader(nranks, train_dataset, batch_size, num_workers, collate_fn) and returns: train_loader