code: def train_cifar(config, data_dir=None):
      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
      )
  
      checkpoint = train.get_checkpoint()
  
      if checkpoint:
          checkpoint_dir = checkpoint.to_directory()
          checkpoint_path = os.path.join(
              checkpoint_dir, "checkpoint.pt"
          )
          checkpoint_state = torch.load(checkpoint_path)
          start_epoch = checkpoint_state["epoch"]
          net.load_state_dict(checkpoint_state["net_state_dict"])
          optimizer.load_state_dict(
              checkpoint_state["optimizer_state_dict"]
          )
      else:
          start_epoch = 0
  
      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=8,
      )
      valloader = torch.utils.data.DataLoader(
          val_subset,
          batch_size=int(config["batch_size"]),
          shuffle=True,
          num_workers=8,
      )
  
      for epoch in range(
          start_epoch, 10
      ):  # loop over the dataset multiple times
          running_loss = 0.0
          epoch_steps = 0
          for i, data in enumerate(trainloader, 0):
              # 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
  
          checkpoint_data = {
              "epoch": epoch,
              "net_state_dict": net.state_dict(),
              "optimizer_state_dict": optimizer.state_dict(),
          }
          with TemporaryDirectory() as tmpdir:
              torch.save(
                  checkpoint_data, os.path.join(tmpdir, "checkpoint.pt")
              )
              train.report(
                  {
                      "loss": val_loss / val_steps,
                      "accuracy": correct / total,
                  },
                  checkpoint=Checkpoint.from_directory(tmpdir),
              )
      print("Finished Training")
  
imports: from functools import partial
  import os
  from tempfile import TemporaryDirectory
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torch.utils.data import random_split
  import torchvision
  import torchvision.transforms
  from ray import train, tune
  from ray.train import Checkpoint
  from ray.tune.schedulers import ASHAScheduler
label_desc: Use image node_id 2 for calling a global function with example usage: train_cifar(config, data_dir) without return types