code: def run_tfe_benchmarks(
      network="B",
      epochs=5,
      start_epoch=0,
      batch_size=256,
      lr=0.01,
      momentum=0.9,
      weight_decay=1e-6,
      print_freq=10,
      resume="",
      evaluate=True,
      seed=None,
      skip_plaintext=False,
      save_checkpoint_dir="/tmp/tfe_benchmarks",
      save_modelbest_dir="/tmp/tfe_benchmarks_best",
      context_manager=None,
      mnist_dir=None,
  ):
      crypten.init()
  
      if seed is not None:
          random.seed(seed)
          torch.manual_seed(seed)
  
      # create model
      model = create_benchmark_model(network)
  
      # define loss function (criterion) and optimizer
      criterion = nn.CrossEntropyLoss()
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          lr,
          momentum=momentum,
          weight_decay=weight_decay,
      )
  
      # optionally resume from a checkpoint
      best_prec1 = 0
      if resume:
          if os.path.isfile(resume):
              logging.info("=> loading checkpoint '{}'".format(resume))
              checkpoint = torch.load(resume)
              start_epoch = checkpoint["epoch"]
              best_prec1 = checkpoint["best_prec1"]
              model.load_state_dict(checkpoint["state_dict"])
              optimizer.load_state_dict(checkpoint["optimizer"])
              logging.info(
                  "=> loaded checkpoint '{}' (epoch {})".format(
                      resume, checkpoint["epoch"]
                  )
              )
          else:
              logging.info(
                  "=> no checkpoint found at '{}'".format(resume)
              )
  
      # Loading MNIST. Normalizing per pytorch/examples/blob/master/mnist/main.py
      def preprocess_data(context_manager, data_dirname):
          if mnist_dir is not None:
              process_mnist_files(
                  mnist_dir,
                  os.path.join(data_dirname, "MNIST", "processed"),
              )
              download = False
          else:
              download = True
  
          with context_manager:
              if not evaluate:
                  mnist_train = datasets.MNIST(
                      data_dirname,
                      download=download,
                      train=True,
                      transform=transforms.Compose(
                          [
                              transforms.ToTensor(),
                              transforms.Normalize(
                                  (0.1307,), (0.3081,)
                              ),
                          ]
                      ),
                  )
  
              mnist_test = datasets.MNIST(
                  data_dirname,
                  download=download,
                  train=False,
                  transform=transforms.Compose(
                      [
                          transforms.ToTensor(),
                          transforms.Normalize((0.1307,), (0.3081,)),
                      ]
                  ),
              )
          train_loader = (
              torch.utils.data.DataLoader(
                  mnist_train, batch_size=batch_size, shuffle=True
              )
              if not evaluate
              else None
          )
          test_loader = torch.utils.data.DataLoader(
              mnist_test, batch_size=batch_size, shuffle=False
          )
          return train_loader, test_loader
  
      if context_manager is None:
          context_manager = NoopContextManager()
  
      warnings.filterwarnings("ignore")
      data_dir = tempfile.TemporaryDirectory()
      train_loader, val_loader = preprocess_data(
          context_manager, data_dir.name
      )
  
      flatten = False
      if network == "A":
          flatten = True
  
      if evaluate:
          if not skip_plaintext:
              logging.info(
                  "===== Evaluating plaintext benchmark network ====="
              )
              validate(
                  val_loader,
                  model,
                  criterion,
                  print_freq,
                  flatten=flatten,
              )
          private_model = create_private_benchmark_model(
              model, flatten=flatten
          )
          logging.info(
              "===== Evaluating Private benchmark network ====="
          )
          validate(
              val_loader,
              private_model,
              criterion,
              print_freq,
              flatten=flatten,
          )
          # validate_side_by_side(val_loader, model, private_model, flatten=flatten)
          return
  
      os.makedirs(save_checkpoint_dir, exist_ok=True)
      os.makedirs(save_modelbest_dir, exist_ok=True)
  
      for epoch in range(start_epoch, epochs):
          adjust_learning_rate(optimizer, epoch, lr)
  
          # train for one epoch
          train(
              train_loader,
              model,
              criterion,
              optimizer,
              epoch,
              print_freq,
              flatten=flatten,
          )
  
          # evaluate on validation set
          prec1 = validate(
              val_loader, model, criterion, print_freq, flatten=flatten
          )
  
          # remember best prec@1 and save checkpoint
          is_best = prec1 > best_prec1
          best_prec1 = max(prec1, best_prec1)
          checkpoint_file = "checkpoint_bn" + network + ".pth.tar"
          model_best_file = "model_best_bn" + network + ".pth.tar"
          save_checkpoint(
              {
                  "epoch": epoch + 1,
                  "arch": "Benchmark" + network,
                  "state_dict": model.state_dict(),
                  "best_prec1": best_prec1,
                  "optimizer": optimizer.state_dict(),
              },
              is_best,
              filename=os.path.join(
                  save_checkpoint_dir, checkpoint_file
              ),
              model_best=os.path.join(
                  save_modelbest_dir, model_best_file
              ),
          )
      data_dir.cleanup()
      shutil.rmtree(save_checkpoint_dir)
  
imports: import logging
  import os
  import random
  import shutil
  import tempfile
  import time
  import warnings
  import crypten
  import crypten.communicator
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torch.utils.data
  import torch.utils.data.distributed
  from examples.meters import AverageMeter
  from examples.util import NoopContextManager, process_mnist_files
  from torchvision import datasets, transforms
label_desc: Use image node_id 1 for calling a global function with example usage: run_tfe_benchmarks(network, epochs, start_epoch, batch_size, lr, momentum, weight_decay, print_freq, resume, evaluate, seed, skip_plaintext, save_checkpoint_dir, save_modelbest_dir, context_manager, mnist_dir) and returns: train_loader,  test_loader