code: def train(run, rank, model, device, dataset):
      torch.manual_seed(SEED + rank)
  
      dataloader_kwargs = {"batch_size": BATCH_SIZE, "shuffle": True}
      train_loader = torch.utils.data.DataLoader(
          dataset, **dataloader_kwargs
      )
  
      optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
      run.define_metric(f"step_{os.getpid}")
      for epoch in range(1, EPOCHS + 1):
          train_epoch(
              run, epoch, model, device, train_loader, optimizer
          )
  
file_path: wandb/tests/functional_tests/t0_main/torch/t2_mp_simple.py
imports: import os
  import torch
  import torch.multiprocessing
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import wandb
  from PIL import Image
  from torchvision import transforms
label_desc: Use image node_id 1 for calling a global function with example usage: train(run, rank, model, device, dataset) without return types