code: def get_data_loader(model, total_samples, hidden_dim, device):
      batch_size = model.train_micro_batch_size_per_gpu()
      train_data = torch.randn(
          total_samples, hidden_dim, device=device, dtype=torch.half
      )
      train_label = torch.empty(
          total_samples, dtype=torch.long, device=device
      ).random_(hidden_dim)
      train_dataset = torch.utils.data.TensorDataset(
          train_data, train_label
      )
      sampler = DistributedSampler(train_dataset)
      train_loader = torch.utils.data.DataLoader(
          train_dataset, batch_size=batch_size, sampler=sampler
      )
      return train_loader
  
file_path: DeepSpeed/tests/small_model_debugging/partial_offload_test.py
imports: import os
  import json
  import argparse
  import torch
  import deepspeed
  from torch.utils.data.distributed import DistributedSampler
  import deepspeed.comm
label_desc: Use image node_id 2 for calling a global function with example usage: get_data_loader(model, total_samples, hidden_dim, device) and returns: train_loader