code: def train_func_per_worker(config: Dict):
      lr = config["lr"]
      epochs = config["epochs"]
      batch_size = config["batch_size_per_worker"]
  
      # Get dataloaders inside the worker training function
      train_dataloader, test_dataloader = get_dataloaders(
          batch_size=batch_size
      )
  
      # [1] Prepare Dataloader for distributed training
      # Shard the datasets among workers and move batches to the correct device
      # =======================================================================
      train_dataloader = ray.train.torch.prepare_data_loader(
          train_dataloader
      )
      test_dataloader = ray.train.torch.prepare_data_loader(
          test_dataloader
      )
  
      model = NeuralNetwork()
  
      # [2] Prepare and wrap your model with DistributedDataParallel
      # Move the model to the correct GPU/CPU device
      # ============================================================
      model = ray.train.torch.prepare_model(model)
  
      loss_fn = nn.CrossEntropyLoss()
      optimizer = torch.optim.SGD(
          model.parameters(), lr=lr, momentum=0.9
      )
  
      # Model training loop
      for epoch in range(epochs):
          if ray.train.get_context().get_world_size() > 1:
              # Required for the distributed sampler to shuffle properly across epochs.
              train_dataloader.sampler.set_epoch(epoch)
  
          model.train()
          for X, y in tqdm(
              train_dataloader, desc=f"Train Epoch {epoch}"
          ):
              pred = model(X)
              loss = loss_fn(pred, y)
  
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
  
          model.eval()
          test_loss, num_correct, num_total = 0, 0, 0
          with torch.no_grad():
              for X, y in tqdm(
                  test_dataloader, desc=f"Test Epoch {epoch}"
              ):
                  pred = model(X)
                  loss = loss_fn(pred, y)
  
                  test_loss += loss.item()
                  num_total += y.shape[0]
                  num_correct += (pred.argmax(1) == y).sum().item()
  
          test_loss /= len(test_dataloader)
          accuracy = num_correct / num_total
  
          # [3] Report metrics to Ray Train
          # ===============================
          ray.train.report(
              metrics={"loss": test_loss, "accuracy": accuracy}
          )
  
file_path: ray/python/ray/train/examples/pytorch/torch_fashion_mnist_example.py
imports: import os
  from typing import Dict
  import torch
  from filelock import FileLock
  from torch import nn
  from torch.utils.data import DataLoader
  from torchvision import datasets, transforms
  from torchvision.transforms import Normalize, ToTensor
  from tqdm import tqdm
  import ray.train
  from ray.train import ScalingConfig
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 2 for calling a global function with example usage: train_func_per_worker(config) without return types