code: def train_epoch(run, epoch, model, device, data_loader, optimizer):
      model.train()
      pid = os.getpid()
      for batch_idx, (data, target) in enumerate(data_loader):
          optimizer.zero_grad()
          output = model(data.to(device))
          loss = F.nll_loss(output, target.to(device))
          loss.backward()
          optimizer.step()
          if batch_idx % LOG_INTERVAL == 0:
              print(
                  "{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                      pid,
                      epoch,
                      batch_idx * len(data),
                      len(data_loader.dataset),
                      100.0 * batch_idx / len(data_loader),
                      loss.item(),
                  )
              )
              run.log(
                  {
                      "loss": loss.item(),
                      f"step_{pid}": batch_idx * len(data),
                  }
              )
  
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 2 for calling a global function with example usage: train_epoch(run, epoch, model, device, data_loader, optimizer) without return types