code: def load_data(
      split="train",
      pca=None,
      clusters=None,
      bandwidth=1.0,
      download_mnist_func=download_mnist,
  ):
      """
      Loads split from the MNIST dataset and returns data.
      """
  
      # download the MNIST dataset:
      mnist = download_mnist_func(split)
  
      # preprocess the MNIST dataset:
      context = mnist.data.float().div_(255.0)
      context = context.view(context.size(0), -1)
  
      # apply PCA:
      if pca is not None:
          context -= torch.mean(context, dim=0, keepdim=True)
          context = context.matmul(pca)
          context /= torch.norm(context, dim=1, keepdim=True)
  
      # compute rewards (based on clustering if clusters defined, 0-1 otherwise):
      if clusters is not None:
          assert clusters.size(1) == context.size(
              1
          ), "cluster dimensionality does not match data dimensionality"
          rewards = examples.util.kmeans_inference(
              context, clusters, hard=False, bandwidth=bandwidth
          )
      else:
          rewards = examples.util.onehot(mnist.targets.long())
  
      # return data:
      return context, rewards
  
file_path: CrypTen/examples/bandits/launcher.py
imports: import argparse
  import logging
  import os
  import random
  import examples.util
  import torch
  import visdom
  from examples.multiprocess_launcher import MultiProcessLauncher
  from examples.util import NoopContextManager, process_mnist_files
  from torchvision.datasets.mnist import MNIST
label_desc: Use image node_id 3 for calling a global function with example usage: load_data(split, pca, clusters, bandwidth, download_mnist_func) and returns: context,  rewards