code: def train(
      args,
      model,
      train_loader,
      optimizer,
      privacy_engine,
      epoch,
      device,
  ):
      start_time = datetime.now()
  
      model.train()
      criterion = nn.CrossEntropyLoss()
  
      losses = []
      top1_acc = []
  
      if args.grad_sample_mode == "no_op":
          # Functorch prepare
          fmodel, _fparams = make_functional(model)
  
          def compute_loss_stateless_model(params, sample, target):
              batch = sample.unsqueeze(0)
              targets = target.unsqueeze(0)
  
              predictions = fmodel(params, batch)
              loss = criterion(predictions, targets)
              return loss
  
          ft_compute_grad = grad_and_value(compute_loss_stateless_model)
          ft_compute_sample_grad = vmap(
              ft_compute_grad, in_dims=(None, 0, 0)
          )
          # Using model.parameters() instead of fparams
          # as fparams seems to not point to the dynamically updated parameters
          params = list(model.parameters())
  
      for i, (images, target) in enumerate(tqdm(train_loader)):
          images = images.to(device)
          target = target.to(device)
  
          # compute output
          output = model(images)
  
          if args.grad_sample_mode == "no_op":
              (
                  per_sample_grads,
                  per_sample_losses,
              ) = ft_compute_sample_grad(params, images, target)
              per_sample_grads = [g.detach() for g in per_sample_grads]
              loss = torch.mean(per_sample_losses)
              for p, g in zip(params, per_sample_grads):
                  p.grad_sample = g
          else:
              loss = criterion(output, target)
              preds = np.argmax(output.detach().cpu().numpy(), axis=1)
              labels = target.detach().cpu().numpy()
  
              # measure accuracy and record loss
              acc1 = accuracy(preds, labels)
              top1_acc.append(acc1)
  
              # compute gradient and do SGD step
              loss.backward()
  
          losses.append(loss.item())
  
          # make sure we take a step after processing the last mini-batch in the
          # epoch to ensure we start the next epoch with a clean state
          optimizer.step()
          optimizer.zero_grad()
  
          if i % args.print_freq == 0:
              if not args.disable_dp:
                  epsilon = privacy_engine.accountant.get_epsilon(
                      delta=args.delta
                  )
                  print(
                      f"\tTrain Epoch: {epoch} \t"
                      f"Loss: {np.mean(losses):.6f} "
                      f"Acc@1: {np.mean(top1_acc):.6f} "
                      f"(ε = {epsilon:.2f}, δ = {args.delta})"
                  )
              else:
                  print(
                      f"\tTrain Epoch: {epoch} \t"
                      f"Loss: {np.mean(losses):.6f} "
                      f"Acc@1: {np.mean(top1_acc):.6f} "
                  )
      train_duration = datetime.now() - start_time
      return train_duration
  
imports: import argparse
  import logging
  import os
  import shutil
  import sys
  from datetime import datetime, timedelta
  import numpy
  import torch
  import torch.nn
  import torch.optim
  import torch.utils.data
  import torch.utils.data.distributed
  import torchvision.transforms
  from opacus import PrivacyEngine
  from opacus.distributed import DifferentiallyPrivateDistributedDataParallel
  from opacus.grad_sample.functorch import make_functional
  from torch.func import grad, grad_and_value, vmap
  from torch.nn.parallel import DistributedDataParallel
  from torchvision.datasets import CIFAR10
  from tqdm import tqdm
label_desc: Use image node_id 6 for calling a global function with example usage: train(args, model, train_loader, optimizer, privacy_engine, epoch, device) and returns: train_duration, loss