code: def train_mnist(config):
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
      train_loader, test_loader = get_data_loaders()
  
      model = ConvNet()
      model.to(device)
  
      optimizer = optim.SGD(
          model.parameters(),
          lr=config["lr"],
          momentum=config["momentum"],
      )
  
      for _i in range(5):
          train(model, optimizer, train_loader, device=device)
          acc = test(model, test_loader, device=device)
  
          # When using WandbLogger, the metrics reported to tune are also logged in the W&B dashboard
          tune.report(mean_accuracy=acc)
  
          # @wandb_mixin enables logging custom metrics using wandb.log()
          error_rate = 100 * (1 - acc)
          wandb.log({"error_rate": error_rate})
  
file_path: wandb/tests/functional_tests/t0_main/raytune/t5_colab.py
imports: import random
  import numpy
  import torch
  import torch.optim
  import wandb
  from ray import tune
  from ray.tune.examples.mnist_pytorch import ConvNet, get_data_loaders, test, train
  from ray.tune.integration.wandb import WandbLogger, wandb_mixin
label_desc: Use image node_id 1 for calling a global function with example usage: train_mnist(config) without return types