code: def train_func_distributed():
      num_epochs = 3
      batch_size = 64
  
      dataset = get_dataset()
      dataloader = DataLoader(
          dataset, batch_size=batch_size, shuffle=True
      )
      dataloader = ray.train.torch.prepare_data_loader(dataloader)
  
      model = NeuralNetwork()
      model = ray.train.torch.prepare_model(model)
  
      criterion = nn.CrossEntropyLoss()
      optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
  
      for epoch in range(num_epochs):
          if ray.train.get_context().get_world_size() > 1:
              dataloader.sampler.set_epoch(epoch)
  
          for inputs, labels in dataloader:
              optimizer.zero_grad()
              pred = model(inputs)
              loss = criterion(pred, labels)
              loss.backward()
              optimizer.step()
          print(f"epoch: {epoch}, loss: {loss.item()}")
  
file_path: ray/python/ray/train/examples/pytorch/torch_quick_start.py
imports: import torch
  import torch.nn
  from torch.utils.data import DataLoader
  from torchvision import datasets
  from torchvision.transforms import ToTensor
  import ray.train.torch
label_desc: Use image node_id 3 for calling a global function with example usage: train_func_distributed() without return types