code: def train_cifar(config):
      net = Net(config["l1"], config["l2"])
  
      device = "cpu"
      if torch.cuda.is_available():
          device = "cuda:0"
          if torch.cuda.device_count() > 1:
              net = nn.DataParallel(net)
      net.to(device)
  
      criterion = nn.CrossEntropyLoss()
      optimizer = optim.SGD(
          net.parameters(), lr=config["lr"], momentum=0.9
      )
  
      # Load existing checkpoint through `get_checkpoint()` API.
      if train.get_checkpoint():
          loaded_checkpoint = train.get_checkpoint()
          with loaded_checkpoint.as_directory() as loaded_checkpoint_dir:
              model_state, optimizer_state = torch.load(
                  os.path.join(loaded_checkpoint_dir, "checkpoint.pt")
              )
              net.load_state_dict(model_state)
              optimizer.load_state_dict(optimizer_state)
  
      if config["smoke_test"]:
          trainset, testset = load_test_data()
      else:
          trainset, testset = load_data(DATA_DIR)
  
      test_abs = int(len(trainset) * 0.8)
      train_subset, val_subset = random_split(
          trainset, [test_abs, len(trainset) - test_abs]
      )
  
      trainloader = torch.utils.data.DataLoader(
          train_subset,
          batch_size=int(config["batch_size"]),
          shuffle=True,
          num_workers=0 if config["smoke_test"] else 8,
      )
      valloader = torch.utils.data.DataLoader(
          val_subset,
          batch_size=int(config["batch_size"]),
          shuffle=True,
          num_workers=0 if config["smoke_test"] else 8,
      )
  
      for epoch in range(10):  # loop over the dataset multiple times
          running_loss = 0.0
          epoch_steps = 0
          for i, data in enumerate(trainloader):
              # get the inputs; data is a list of [inputs, labels]
              inputs, labels = data
              inputs, labels = inputs.to(device), labels.to(device)
  
              # zero the parameter gradients
              optimizer.zero_grad()
  
              # forward + backward + optimize
              outputs = net(inputs)
              loss = criterion(outputs, labels)
              loss.backward()
              optimizer.step()
  
              # print statistics
              running_loss += loss.item()
              epoch_steps += 1
              if i % 2000 == 1999:  # print every 2000 mini-batches
                  print(
                      "[%d, %5d] loss: %.3f"
                      % (epoch + 1, i + 1, running_loss / epoch_steps)
                  )
                  running_loss = 0.0
  
          # Validation loss
          val_loss = 0.0
          val_steps = 0
          total = 0
          correct = 0
          for i, data in enumerate(valloader, 0):
              with torch.no_grad():
                  inputs, labels = data
                  inputs, labels = inputs.to(device), labels.to(device)
  
                  outputs = net(inputs)
                  _, predicted = torch.max(outputs.data, 1)
                  total += labels.size(0)
                  correct += (predicted == labels).sum().item()
  
                  loss = criterion(outputs, labels)
                  val_loss += loss.cpu().numpy()
                  val_steps += 1
  
          # Here we save a checkpoint. It is automatically registered with
          # Ray Tune and will potentially be accessed through in ``get_checkpoint()``
          # in future iterations.
          # Note to save a file like checkpoint, you still need to put it under a directory
          # to construct a checkpoint.
          with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
              path = os.path.join(temp_checkpoint_dir, "checkpoint.pt")
              torch.save(
                  (net.state_dict(), optimizer.state_dict()), path
              )
              checkpoint = Checkpoint.from_directory(
                  temp_checkpoint_dir
              )
              train.report(
                  {
                      "loss": (val_loss / val_steps),
                      "accuracy": correct / total,
                  },
                  checkpoint=checkpoint,
              )
      print("Finished Training")
  
imports: import numpy
  import os
  import tempfile
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from filelock import FileLock
  from torch.utils.data import random_split
  import torchvision
  import torchvision.transforms
  from typing import Dict
  import ray
  from ray import train, tune
  from ray.train import Checkpoint
  from ray.tune.schedulers import ASHAScheduler
label_desc: Use image node_id 3 for calling a global function with example usage: train_cifar(config) without return types