text: import argparse
  import logging
  import os
  import random
  import socket
  import sys
  import traceback
  
  import numpy as np
  import psutil
  import setproctitle
  import torch
  
  from torch import nn
  from torch.utils.data import DataLoader
  from torchvision import datasets, transforms
  
  # add the FedML root directory to the python path
  sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "./../../../../")))
  sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "./../../../")))
  sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "")))
  
  ################# Testing for model_dimension ########################
  from fedml_api.distributed.lightsecagg.utils import model_dimension
  
  #model_dimension is to get the the dimension for each layer of model and the total dimension of all the layers
  
  print(123)