code: def per_device_launch_fn(current_gpu_index, num_gpu):
      # Setup the process groups
      setup_device(current_gpu_index, num_gpu)
  
      dataset = get_dataset()
      model = get_model()
  
      # prepare the dataloader
      dataloader = prepare_dataloader(
          dataset, current_gpu_index, num_gpu, batch_size
      )
  
      # Instantiate the torch optimizer
      optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
  
      # Instantiate the torch loss function
      loss_fn = torch.nn.CrossEntropyLoss()
  
      # Put model on device
      model = model.to(current_gpu_index)
      ddp_model = torch.nn.parallel.DistributedDataParallel(
          model,
          device_ids=[current_gpu_index],
          output_device=current_gpu_index,
      )
  
      train_model(ddp_model, dataloader, num_epochs, optimizer, loss_fn)
  
      cleanup()
  
file_path: keras/guides/distributed_training_with_torch.py
imports: import os
  import torch
  import numpy
  import keras
label_desc: Use image node_id 7 for calling a global function with example usage: per_device_launch_fn(current_gpu_index, num_gpu) without return types