code: def train_func(config):
      num_epochs = config["num_epochs"]
      batch_size = config["batch_size"]
      num_layers = config["num_layers"]
      num_hidden = config["num_hidden"]
      dropout_every = config["dropout_every"]
      dropout_prob = config["dropout_prob"]
      num_features = config["num_features"]
  
      print("Defining model, loss, and optimizer...")
  
      # Setup device.
      device = train.torch.get_device()
      print(f"Device: {device}")
  
      # Setup data.
      train_dataset_iterator = train.get_dataset_shard("train")
      test_dataset_iterator = train.get_dataset_shard("test")
  
      def to_torch_dataset(torch_batch_iterator):
          for batch in torch_batch_iterator:
              label_column = "label"
              labels = batch[label_column].unsqueeze(1)
              features = [
                  batch[col_name].unsqueeze(1)
                  for col_name in batch
                  if col_name != label_column
              ]
              inputs = torch.cat(features, dim=1)
              yield inputs, labels
  
      net = Net(
          n_layers=num_layers,
          n_features=num_features,
          num_hidden=num_hidden,
          dropout_every=dropout_every,
          drop_prob=dropout_prob,
      ).to(device)
      print(net.parameters)
  
      net = train.torch.prepare_model(net)
  
      criterion = nn.BCEWithLogitsLoss()
      optimizer = optim.Adam(net.parameters(), weight_decay=0.0001)
  
      print("Starting training...")
      for epoch in range(num_epochs):
          train_torch_dataset = to_torch_dataset(
              train_dataset_iterator.iter_torch_batches(
                  batch_size=batch_size
              )
          )
          (
              train_running_loss,
              train_num_correct,
              train_num_total,
          ) = train_epoch(
              train_torch_dataset, net, device, criterion, optimizer
          )
          train_acc = train_num_correct / train_num_total
          print(
              f"epoch [{epoch + 1}]: training accuracy: "
              f"{train_num_correct} / {train_num_total} = {train_acc:.4f}"
          )
  
          test_torch_dataset = to_torch_dataset(
              test_dataset_iterator.iter_torch_batches(
                  batch_size=batch_size, drop_last=True
              )
          )
          (
              test_running_loss,
              test_num_correct,
              test_num_total,
          ) = test_epoch(test_torch_dataset, net, device, criterion)
          test_acc = test_num_correct / test_num_total
          print(
              f"epoch [{epoch + 1}]: testing accuracy: "
              f"{test_num_correct} / {test_num_total} = {test_acc:.4f}"
          )
  
          # Checkpoint model.
          with TemporaryDirectory() as tmpdir:
              torch.save(
                  net.module.state_dict(),
                  os.path.join(tmpdir, "checkpoint.pt"),
              )
  
              # Record and log stats.
              print(
                  f"train report on {train.get_context().get_world_rank()}"
              )
              train.report(
                  dict(
                      train_acc=train_acc,
                      train_loss=train_running_loss,
                      test_acc=test_acc,
                      test_loss=test_running_loss,
                  ),
                  checkpoint=Checkpoint.from_directory(tmpdir),
              )
  
imports: import argparse
  import collections
  import os
  import sys
  import time
  from tempfile import TemporaryDirectory
  from typing import Tuple
  import boto3
  import mlflow
  import pandas
  import torch
  import torch.nn
  import torch.optim
  import ray
  from ray import train
  from ray.air.integrations.mlflow import MLflowLoggerCallback
  from ray.data.aggregate import Mean, Std
  from ray.train import Checkpoint, DataConfig, RunConfig, ScalingConfig
  from ray.train.torch.torch_trainer import TorchTrainer
label_desc: Use image node_id 6 for calling a global function with example usage: train_func(config) without return types