code: def train_torch(
      model,
      X,
      y,
      batch_size,
      device,
      lr,
      max_epochs,
  ):
      model = to_device(model, device)
  
      idx_train, idx_valid = next(
          iter(StratifiedKFold(5).split(np.arange(len(X)), y))
      )
      X_train, X_valid, y_train, y_valid = (
          X[idx_train],
          X[idx_valid],
          y[idx_train],
          y[idx_valid],
      )
      dataset_train = torch.utils.data.TensorDataset(
          torch.tensor(X_train),
          torch.tensor(y_train),
      )
      iterator_train = torch.utils.data.DataLoader(
          dataset_train, batch_size=batch_size
      )
      dataset_valid = torch.utils.data.TensorDataset(
          torch.tensor(X_valid),
          torch.tensor(y_valid),
      )
      iterator_valid = torch.utils.data.DataLoader(
          dataset_valid, batch_size=batch_size
      )
  
      optimizer = torch.optim.Adadelta(model.parameters(), lr=lr)
      criterion = nn.NLLLoss()
  
      for epoch in range(max_epochs):
          train_out = train_step(
              model,
              iterator_train,
              batch_size=batch_size,
              device=device,
              criterion=criterion,
              optimizer=optimizer,
          )
          report(y=y_train, epoch=epoch, training=True, **train_out)
  
          valid_out = valid_step(
              model,
              iterator_valid,
              batch_size=batch_size,
              device=device,
              criterion=criterion,
          )
          report(y=y_valid, epoch=epoch, training=False, **valid_out)
  
          print("-" * 50)
  
      return model
  
imports: import argparse
  import time
  import numpy
  from sklearn.datasets import fetch_openml
  from sklearn.model_selection import train_test_split
  from sklearn.model_selection import StratifiedKFold
  from sklearn.metrics import accuracy_score
  from sklearn.utils import shuffle
  import torch
  from torch import nn
  from skorch.utils import to_device
  from skorch import NeuralNetClassifier
  from skorch.callbacks import EpochScoring
label_desc: Use image node_id 4 for calling a global function with example usage: train_torch(model, X, y, batch_size, device, lr, max_epochs) and returns: model