code: def train_func(use_ray: bool, config: Dict):
      local_start_time = time.monotonic()
  
      if use_ray:
          import ray.train as train
  
      batch_size = config["batch_size"]
      lr = config["lr"]
      epochs = config["epochs"]
      shuffle = config.get("shuffle", False)
  
      if use_ray:
          world_size = train.get_context().get_world_size()
          local_rank = distributed.get_rank()
      else:
          world_size = distributed.get_world_size()
          local_rank = distributed.get_rank()
  
      worker_batch_size = batch_size // world_size
  
      # Load datasets. Use download=False to catch errors in preparation, as the
      # data should have already been downloaded.
      training_data = datasets.FashionMNIST(
          root="/tmp/data_fashion_mnist",
          train=True,
          download=False,
          transform=ToTensor(),
      )
  
      test_data = datasets.FashionMNIST(
          root="/tmp/data_fashion_mnist",
          train=False,
          download=False,
          transform=ToTensor(),
      )
  
      if use_ray:
          # Ray adds DistributedSampler in train.torch.prepare_data_loader below
          training_sampler = None
          test_sampler = None
      else:
          # In vanilla PyTorch we create the distributed sampler here
          training_sampler = DistributedSampler(
              training_data, shuffle=shuffle
          )
          test_sampler = DistributedSampler(test_data, shuffle=shuffle)
  
      if not use_ray and config.get("use_gpu", False):
          assert torch.cuda.is_available(), "No GPUs available"
          gpu_id = config.get("gpu_id", 0)
          vanilla_device = torch.device(f"cuda:{gpu_id}")
          torch.cuda.set_device(vanilla_device)
  
          print(
              "Setting GPU ID to",
              gpu_id,
              "with visible devices",
              os.environ.get("CUDA_VISIBLE_DEVICES"),
          )
  
          def collate_fn(x):
              return tuple(
                  x_.to(vanilla_device) for x_ in default_collate(x)
              )
  
      else:
          vanilla_device = torch.device("cpu")
          collate_fn = None
  
      # Create data loaders and potentially pass distributed sampler
      train_dataloader = DataLoader(
          training_data,
          shuffle=shuffle,
          batch_size=worker_batch_size,
          sampler=training_sampler,
          collate_fn=collate_fn,
      )
      test_dataloader = DataLoader(
          test_data,
          shuffle=shuffle,
          batch_size=worker_batch_size,
          sampler=test_sampler,
          collate_fn=collate_fn,
      )
  
      if use_ray:
          # In Ray, we now retrofit the DistributedSampler
          train_dataloader = train.torch.prepare_data_loader(
              train_dataloader
          )
          test_dataloader = train.torch.prepare_data_loader(
              test_dataloader
          )
  
      # Create model.
      model = NeuralNetwork()
  
      # Prepare model
      if use_ray:
          model = train.torch.prepare_model(model)
      else:
          model = model.to(vanilla_device)
  
          if config.get("use_gpu", False):
              model = nn.parallel.DistributedDataParallel(
                  model, device_ids=[gpu_id], output_device=gpu_id
              )
          else:
              model = nn.parallel.DistributedDataParallel(model)
  
      loss_fn = nn.CrossEntropyLoss()
      optimizer = torch.optim.SGD(model.parameters(), lr=lr)
  
      for epoch in range(epochs):
          if world_size > 1:
              train_dataloader.sampler.set_epoch(epoch)
  
          train_epoch(
              train_dataloader,
              model,
              loss_fn,
              optimizer,
              world_size=world_size,
              local_rank=local_rank,
          )
          loss = validate_epoch(
              test_dataloader,
              model,
              loss_fn,
              world_size=world_size,
              local_rank=local_rank,
          )
  
          local_time_taken = time.monotonic() - local_start_time
  
          if use_ray:
              train.report(
                  dict(loss=loss, local_time_taken=local_time_taken)
              )
          else:
              print(f"Reporting loss: {loss:.4f}")
              if local_rank == 0:
                  with open(VANILLA_RESULT_JSON, "w") as f:
                      json.dump(
                          {
                              "loss": loss,
                              "local_time_taken": local_time_taken,
                          },
                          f,
                      )
  
imports: import json
  import os
  import time
  from pathlib import Path
  from typing import Dict, Tuple
  import click
  import numpy
  import torch
  from torch import nn, distributed
  from torch.utils.data import DataLoader, DistributedSampler
  from torch.utils.data.dataloader import default_collate
  from torchvision import datasets
  from torchvision.transforms import ToTensor
label_desc: Use image node_id 4 for calling a global function with example usage: train_func(use_ray, config) and returns: tuple