code: def get_dataset():
      # Load the data and split it between train and test sets
      (x_train, y_train), (
          x_test,
          y_test,
      ) = keras.datasets.mnist.load_data()
  
      # Scale images to the [0, 1] range
      x_train = x_train.astype("float32")
      x_test = x_test.astype("float32")
      # Make sure images have shape (28, 28, 1)
      x_train = np.expand_dims(x_train, -1)
      x_test = np.expand_dims(x_test, -1)
      print("x_train shape:", x_train.shape)
  
      # Create a TensorDataset
      dataset = torch.utils.data.TensorDataset(
          torch.from_numpy(x_train), torch.from_numpy(y_train)
      )
      return dataset
  
file_path: keras/guides/distributed_training_with_torch.py
imports: import os
  import torch
  import numpy
  import keras
label_desc: Use image node_id 2 for calling a global function with example usage: get_dataset() and returns: dataset