code: def get_data(num_samples):
      mnist = fetch_openml("mnist_784", as_frame=False)
      torch.manual_seed(0)
      X = mnist.data.astype("float32").reshape(-1, 1, 28, 28)
      y = mnist.target.astype("int64")
      X, y = shuffle(X, y, random_state=0)
      X, y = X[:num_samples], y[:num_samples]
      X_train, X_test, y_train, y_test = train_test_split(
          X, y, stratify=y, random_state=0
      )
      X_train /= 255
      X_test /= 255
      return X_train, X_test, y_train, y_test
  
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 1 for calling a global function with example usage: get_data(num_samples) and returns: X_train,  X_test,  y_train,  y_test