code: def get_data():
      # 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") / 255
      x_test = x_test.astype("float32") / 255
      # 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)
      print(x_train.shape[0], "train samples")
      print(x_test.shape[0], "test samples")
  
      # Create a TensorDataset
      dataset = torch.utils.data.TensorDataset(
          torch.from_numpy(x_train), torch.from_numpy(y_train)
      )
      return dataset
  
imports: import os
  import torch
  import torch.nn
  import torch.optim
  from keras import layers
  import keras
  import numpy
  import torch.multiprocessing
  import torch.distributed
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader, Dataset
  from torch.utils.data.distributed import DistributedSampler
label_desc: Use image node_id 1 for calling a global function with example usage: get_data() and returns: dataset