diff --git a/CMakeLists.txt b/CMakeLists.txt
old mode 100755
new mode 100644
index dfe2959..6101ef5
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -6,6 +6,7 @@ set(protobuf_MODULE_COMPATIBLE TRUE)
 find_package(Protobuf CONFIG REQUIRED)
 set(protobuf_MODULE_COMPATIBLE TRUE)
 find_package(Protobuf CONFIG REQUIRED)
+find_package(Torch REQUIRED)
 
 set(CMAKE_MODULE_PATH
   "${CMAKE_CURRENT_SOURCE_DIR}/cmake"
@@ -13,10 +14,15 @@ set(CMAKE_MODULE_PATH
   ${CMAKE_MODULE_PATH}
   )
 message("CMAKE Module path: ${CMAKE_MODULE_PATH}")
+
+
 set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib CACHE PATH "Output directory for static libraries")
 
 include_directories(${Protobuf_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}/include)
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-rtti -fPIC")
+
+# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-rtti -fPIC")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fPIC")
+
 set (CMAKE_CXX_STANDARD 17)
 
 option(LLVM_MLBRIDGE "MLCompilerBridge install for LLVM" OFF)
@@ -65,10 +71,12 @@ if(LLVM_MLBRIDGE)
   LINK_LIBS
   ModelRunnerLib
   SerDesLib
-  )
+  )  
+
+  target_link_libraries(LLVMMLBridge PRIVATE "${TORCH_LIBRARIES}")
 
-  target_include_directories(LLVMMLBridge SYSTEM PUBLIC ${Protobuf_INCLUDE_DIRS} ${TENSORFLOW_AOT_PATH}/include)
-  target_include_directories(LLVMMLBridge PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
+  target_include_directories(LLVMMLBridge SYSTEM PRIVATE ${Protobuf_INCLUDE_DIRS} ${TENSORFLOW_AOT_PATH}/include)
+  target_include_directories(LLVMMLBridge PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)
   install(TARGETS LLVMMLBridge DESTINATION lib)
   add_custom_command(TARGET LLVMMLBridge
     POST_BUILD
@@ -79,28 +87,35 @@ else()
   llvm_map_components_to_libnames(llvm_libs support core irreader analysis TransformUtils)
 
   add_library(MLCompilerBridge STATIC tools.cpp)
-  target_link_libraries(MLCompilerBridge PUBLIC SerDesLib ModelRunnerLib ONNXModelRunnerLib ${llvm_libs})
+  target_link_libraries(MLCompilerBridge PRIVATE SerDesLib ModelRunnerLib ONNXModelRunnerLib ${llvm_libs})
+  target_link_libraries(MLCompilerBridge PRIVATE "${TORCH_LIBRARIES}")
   set_target_properties(MLCompilerBridge PROPERTIES ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
   set_property(TARGET MLCompilerBridge PROPERTY POSITION_INDEPENDENT_CODE 1)
   install(TARGETS MLCompilerBridge DESTINATION lib)
 
 
   add_library(MLCompilerBridgeC STATIC $<TARGET_OBJECTS:ModelRunnerCWrapper>)
-  target_link_libraries(MLCompilerBridgeC PUBLIC SerDesCLib ModelRunnerCLib ONNXModelRunnerLib ${llvm_libs})
-  target_include_directories(MLCompilerBridgeC PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${LLVM_INCLUDE_DIRS})
+  target_link_libraries(MLCompilerBridgeC PRIVATE SerDesCLib ModelRunnerCLib ONNXModelRunnerLib ${llvm_libs})
+
+  target_include_directories(MLCompilerBridgeC PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include ${LLVM_INCLUDE_DIRS})
   target_compile_features(MLCompilerBridgeC PRIVATE cxx_std_17)
   target_compile_definitions(MLCompilerBridgeC PRIVATE C_LIBRARY)
   set_property(TARGET MLCompilerBridgeC PROPERTY POSITION_INDEPENDENT_CODE 1)
   set_target_properties(MLCompilerBridgeC PROPERTIES ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
   install(TARGETS MLCompilerBridgeC DESTINATION lib)
   if(MLBRIDGE_ENABLE_TEST)
-    add_executable(MLCompilerBridgeTest $<TARGET_OBJECTS:MLBridgeCPPTest>)
-    target_link_libraries(MLCompilerBridgeTest PUBLIC MLCompilerBridge)
-    set_target_properties(MLCompilerBridgeTest PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
-    install(TARGETS MLCompilerBridgeTest DESTINATION bin)
-  endif()
-endif(LLVM_MLBRIDGE)
+  add_executable(MLCompilerBridgeTest $<TARGET_OBJECTS:MLBridgeCPPTest>)
+  target_link_libraries(MLCompilerBridgeTest PRIVATE MLCompilerBridge)
+  set_target_properties(MLCompilerBridgeTest PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
+  install(TARGETS MLCompilerBridgeTest DESTINATION bin)
+endif()
+endif(LLVM_MLBRIDGE) 
 
 install(DIRECTORY include/ DESTINATION include)
 install(DIRECTORY CompilerInterface DESTINATION MLModelRunner/CompilerInterface)
 file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/CompilerInterface DESTINATION ${CMAKE_BINARY_DIR}/MLModelRunner/)
+set_property(TARGET MLCompilerBridge PROPERTY CXX_STANDARD 17)
+
+
+
+message(STATUS "TensorFlow HH version: ${TensorFlow_VERSION}")
\ No newline at end of file
diff --git a/MLModelRunner/C/CMakeLists.txt b/MLModelRunner/C/CMakeLists.txt
index bcfa357..d0903f7 100644
--- a/MLModelRunner/C/CMakeLists.txt
+++ b/MLModelRunner/C/CMakeLists.txt
@@ -1,2 +1,4 @@
+find_package(Torch REQUIRED)
 add_library(ModelRunnerCWrapper OBJECT PipeModelRunnerCWrapper.cpp ONNXModelRunnerCWrapper.cpp)
-target_include_directories(ModelRunnerCWrapper PUBLIC ${TENSORFLOW_AOT_PATH}/include)
+target_include_directories(ModelRunnerCWrapper PRIVATE ${TENSORFLOW_AOT_PATH}/include)
+target_link_libraries(ModelRunnerCWrapper PRIVATE  "${TORCH_LIBRARIES}")
diff --git a/MLModelRunner/CMakeLists.txt b/MLModelRunner/CMakeLists.txt
old mode 100755
new mode 100644
index 57a1cde..efe3a1c
--- a/MLModelRunner/CMakeLists.txt
+++ b/MLModelRunner/CMakeLists.txt
@@ -1,3 +1,4 @@
+find_package(Torch REQUIRED)  
 add_subdirectory(Utils)
 set(PROTOS_DIRECTORY "" CACHE PATH "Path to the directory containing the proto files")
 if(NOT PROTOS_DIRECTORY STREQUAL "")
@@ -35,7 +36,7 @@ if (NOT TENSORFLOW_AOT_PATH STREQUAL "")
 endif()
 
 add_library(ModelRunnerCLib OBJECT PipeModelRunner.cpp)
-target_link_libraries(ModelRunnerCLib PUBLIC ModelRunnerCUtils ONNXModelRunnerLib)
+target_link_libraries(ModelRunnerCLib PRIVATE ModelRunnerCUtils ONNXModelRunnerLib)
 target_compile_definitions(ModelRunnerCLib PRIVATE C_LIBRARY)
 
 if(LLVM_MLBRIDGE)
@@ -44,9 +45,11 @@ else()
   add_library(ModelRunnerLib OBJECT PipeModelRunner.cpp)
 endif(LLVM_MLBRIDGE)
 
-target_link_libraries(ModelRunnerLib PUBLIC ModelRunnerUtils ONNXModelRunnerLib)
+target_link_libraries(ModelRunnerLib PRIVATE ModelRunnerUtils ONNXModelRunnerLib "${TORCH_LIBRARIES}")
+
 
 if(NOT PROTOS_DIRECTORY STREQUAL "")
-  target_link_libraries(ModelRunnerLib PUBLIC gRPCModelRunnerLib)
+  target_link_libraries(ModelRunnerLib PRIVATE gRPCModelRunnerLib)
 endif()
 set_property(TARGET ModelRunnerLib PROPERTY POSITION_INDEPENDENT_CODE 1)
+set_property(TARGET ModelRunnerLib PROPERTY CXX_STANDARD 17)
diff --git a/MLModelRunner/ONNXModelRunner/CMakeLists.txt b/MLModelRunner/ONNXModelRunner/CMakeLists.txt
old mode 100755
new mode 100644
index d8f197c..a96c57e
--- a/MLModelRunner/ONNXModelRunner/CMakeLists.txt
+++ b/MLModelRunner/ONNXModelRunner/CMakeLists.txt
@@ -1,4 +1,4 @@
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -DORT_NO_EXCEPTIONS")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}  -Wall -Wextra -DORT_NO_EXCEPTIONS")
 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Wextra")
 
 find_package(Onnxruntime REQUIRED)
@@ -13,30 +13,30 @@ option(onnxruntime_USE_TENSORRT "Build with TensorRT support" OFF)
 
 set(ONNXRUNTIME_ROOTDIR "" CACHE PATH "Directory that contains ONNXRuntime" )
 if(NOT ONNXRUNTIME_ROOTDIR)
-  message( FATAL_ERROR "Set path to Onnx runtime in -DONNXRUNTIME_ROOTDIR variable" )
+message( FATAL_ERROR "Set path to Onnx runtime in -DONNXRUNTIME_ROOTDIR variable" )
 endif()
 
 if(onnxruntime_USE_CUDA)
-  add_definitions(-DUSE_CUDA)
+add_definitions(-DUSE_CUDA)
 endif()
 if(onnxruntime_USE_OPENVINO)
-  add_definitions(-DUSE_OPENVINO)
+add_definitions(-DUSE_OPENVINO)
 endif()
 if(onnxruntime_USE_NNAPI_BUILTIN)
-  add_definitions(-DUSE_NNAPI)
+add_definitions(-DUSE_NNAPI)
 endif()
 if(onnxruntime_USE_DNNL)
-  add_definitions(-DUSE_DNNL)
+add_definitions(-DUSE_DNNL)
 endif()
 if(onnxruntime_USE_NUPHAR)
-  add_definitions(-DUSE_NUPHAR)
+add_definitions(-DUSE_NUPHAR)
 endif()
 if(onnxruntime_USE_TENSORRT)
-  add_definitions(-DUSE_TENSORRT)
+add_definitions(-DUSE_TENSORRT)
 endif()
 if(onnxruntime_USE_DML)
-  message("Enabling DML")
-  add_definitions(-DUSE_DML)
+message("Enabling DML")
+add_definitions(-DUSE_DML)
 endif()
 
 if(LLVM_MLBRIDGE)
@@ -47,10 +47,13 @@ ONNXModelRunner.cpp
 )
 else()
 add_library(ONNXModelRunnerLib OBJECT onnx.cpp
-      agent.cpp
-      ONNXModelRunner.cpp
-    )
+agent.cpp
+ONNXModelRunner.cpp
+)
 endif(LLVM_MLBRIDGE)
 
 target_link_libraries(ONNXModelRunnerLib PRIVATE Onnxruntime::Onnxruntime)
-target_include_directories(ONNXModelRunnerLib PUBLIC ${TENSORFLOW_AOT_PATH}/include)
+target_include_directories(ONNXModelRunnerLib PRIVATE ${TENSORFLOW_AOT_PATH}/include)
+find_package(Torch REQUIRED)
+target_link_libraries(ONNXModelRunnerLib PRIVATE ${TORCH_LIBRARIES})
+
diff --git a/MLModelRunner/ONNXModelRunner/ONNXModelRunner.cpp b/MLModelRunner/ONNXModelRunner/ONNXModelRunner.cpp
old mode 100755
new mode 100644
diff --git a/MLModelRunner/ONNXModelRunner/agent.cpp b/MLModelRunner/ONNXModelRunner/agent.cpp
old mode 100755
new mode 100644
diff --git a/MLModelRunner/PTModelRunner.cpp b/MLModelRunner/PTModelRunner.cpp
new file mode 100644
index 0000000..96f7939
--- /dev/null
+++ b/MLModelRunner/PTModelRunner.cpp
@@ -0,0 +1,7 @@
+//=== PTModelRunner.cpp - PTModelRunner Implementation ---*- C++ -*-===//
+//
+// Part of the MLCompilerBridge Project
+//
+//===------------------===//
+
+#include "MLModelRunner/PTModelRunner.h"
diff --git a/MLModelRunner/Utils/CMakeLists.txt b/MLModelRunner/Utils/CMakeLists.txt
old mode 100755
new mode 100644
index 0ee5d0d..2844543
--- a/MLModelRunner/Utils/CMakeLists.txt
+++ b/MLModelRunner/Utils/CMakeLists.txt
@@ -2,7 +2,7 @@ if(LLVM_MLBRIDGE)
   add_llvm_component_library(ModelRunnerUtils
   MLConfig.cpp
   )
-  target_include_directories(ModelRunnerUtils PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/../../include)
+  target_include_directories(ModelRunnerUtils PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../include)
 
 else()
   add_library(ModelRunnerCUtils OBJECT MLConfig.cpp)
diff --git a/MLModelRunner/gRPCModelRunner/CMakeLists.txt b/MLModelRunner/gRPCModelRunner/CMakeLists.txt
old mode 100755
new mode 100644
index 877afef..3687585
--- a/MLModelRunner/gRPCModelRunner/CMakeLists.txt
+++ b/MLModelRunner/gRPCModelRunner/CMakeLists.txt
@@ -1,5 +1,6 @@
 set(protobuf_MODULE_COMPATIBLE TRUE)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+find_package(Torch REQUIRED)
 find_package(Protobuf CONFIG REQUIRED)
 set(_PROTOBUF_LIBPROTOBUF protobuf::libprotobuf)
 set(_REFLECTION gRPC::grpc++_reflection)
@@ -12,7 +13,7 @@ endif()
 # Find gRPC installation
 # Looks for gRPCConfig.cmake file installed by gRPC's cmake installation.
 
-find_package(gRPC 1.58.0 EXACT CONFIG REQUIRED)
+find_package(gRPC 1.34.0 EXACT CONFIG REQUIRED)
 message(STATUS "Using gRPC ${gRPC_VERSION}")
 
 set(_GRPC_GRPCPP gRPC::grpc++)
@@ -96,9 +97,11 @@ target_link_libraries(gRPCModelRunnerLib
 	  PRIVATE ${_REFLECTION}
 	          ${_GRPC_GRPCPP}
 	          ${_PROTOBUF_LIBPROTOBUF}
+            "${TORCH_LIBRARIES}"
 )
 
 target_include_directories(gRPCModelRunnerLib
   PRIVATE ${CMAKE_BINARY_DIR}/include/Service
-  PUBLIC ${Protobuf_INCLUDE_DIRS}
+  PRIVATE ${Protobuf_INCLUDE_DIRS}
 )
+set_property(TARGET gRPCModelRunnerLib PROPERTY CXX_STANDARD 17)
diff --git a/SerDes/CMakeLists.txt b/SerDes/CMakeLists.txt
old mode 100755
new mode 100644
index 9125409..59deef3
--- a/SerDes/CMakeLists.txt
+++ b/SerDes/CMakeLists.txt
@@ -1,19 +1,28 @@
 set(protobuf_MODULE_COMPATIBLE TRUE)
 find_package(Protobuf CONFIG REQUIRED)
+set(protobuf_MODULE_COMPATIBLE TRUE)
+find_package(Protobuf CONFIG REQUIRED)
+
 
 if(LLVM_MLBRIDGE)
-  add_llvm_library(SerDesLib
-  TensorSpec.cpp
-  jsonSerDes.cpp
-  bitstreamSerDes.cpp
-  protobufSerDes.cpp
-  tensorflowSerDes.cpp
-  )
-  target_compile_definitions(SerDesLib PRIVATE LLVM_MLBRIDGE)
+add_llvm_library(SerDesLib
+TensorSpec.cpp
+jsonSerDes.cpp
+bitstreamSerDes.cpp
+protobufSerDes.cpp
+pytorchSerDes.cpp
+tensorflowSerDes.cpp
+)
+target_compile_definitions(SerDesLib PRIVATE LLVM_MLBRIDGE)
 else()
-  add_library(SerDesLib OBJECT TensorSpec.cpp jsonSerDes.cpp bitstreamSerDes.cpp protobufSerDes.cpp tensorflowSerDes.cpp)
+# add_library(SerDesLib OBJECT TensorSpec.cpp jsonSerDes.cpp bitstreamSerDes.cpp protobufSerDes.cpp tensorflowSerDes.cpp pytorchSerDes.cpp)
+add_library(SerDesLib OBJECT TensorSpec.cpp jsonSerDes.cpp bitstreamSerDes.cpp pytorchSerDes.cpp)
+
 
-  add_library(SerDesCLib OBJECT TensorSpec.cpp jsonSerDes.cpp bitstreamSerDes.cpp)
+add_library(SerDesCLib OBJECT TensorSpec.cpp jsonSerDes.cpp bitstreamSerDes.cpp )
 endif()
-target_include_directories(SerDesLib PUBLIC ${TENSORFLOW_AOT_PATH}/include)
+target_include_directories(SerDesLib PRIVATE ${TENSORFLOW_AOT_PATH}/include)
+# find_package(Torch REQUIRED)
 target_link_libraries(SerDesLib PRIVATE tf_xla_runtime)
+target_link_libraries(SerDesLib PRIVATE ${TORCH_LIBRARIES})
+set_property(TARGET SerDesLib PROPERTY CXX_STANDARD 17)
diff --git a/SerDes/protobufSerDes.cpp b/SerDes/protobufSerDes.cpp
index 1e543f1..9e43e89 100644
--- a/SerDes/protobufSerDes.cpp
+++ b/SerDes/protobufSerDes.cpp
@@ -23,6 +23,8 @@
 #include <type_traits>
 #include <vector>
 
+using namespace google::protobuf;
+
 namespace MLBridge {
 inline void ProtobufSerDes::setFeature(const std::string &name,
                                        const int value) {
diff --git a/SerDes/pytorchSerDes.cpp b/SerDes/pytorchSerDes.cpp
new file mode 100644
index 0000000..eb98839
--- /dev/null
+++ b/SerDes/pytorchSerDes.cpp
@@ -0,0 +1,136 @@
+//=== MLCompilerBridge/SerDes/pytorchSerDes.cpp - Pytorch SerDes Implementation ---*- C++ -*-===//
+//
+// Part of the MLCompilerBridge Project
+//
+//===------------------===//
+
+#include "SerDes/pytorchSerDes.h"
+#include "SerDes/baseSerDes.h"
+
+namespace MLBridge {
+
+void PytorchSerDes::setFeature(const std::string &Name, const int Value) {
+    auto tensor = torch::tensor({Value}, torch::kInt32);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const long Value) {
+    auto tensor = torch::tensor({Value}, torch::kInt64);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const float Value) {
+    auto tensor = torch::tensor({Value}, torch::kFloat32);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const double Value) {
+    auto tensor = torch::tensor({Value}, torch::kFloat64);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const std::string Value) {
+    std::vector<int8_t> encoded_str(Value.begin(), Value.end());
+    auto tensor = torch::tensor(encoded_str, torch::kInt8);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const bool Value) {
+    auto tensor = torch::tensor({Value}, torch::kBool);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const std::vector<int> &Value) {
+    auto tensor = torch::tensor(Value, torch::kInt32);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const std::vector<long> &Value) {
+    auto tensor = torch::tensor(Value, torch::kInt64);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const std::vector<float> &Value) {
+    auto tensor = torch::tensor(Value, torch::kFloat32);
+    tensor = tensor.reshape({1, Value.size()});
+    inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const std::vector<double> &Value) {
+    auto tensor = torch::tensor(Value, torch::kFloat64);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const std::vector<std::string> &Value) {
+    std::vector<int8_t> flat_vec;
+    for (const auto& str : Value) {
+        flat_vec.insert(flat_vec.end(), str.begin(), str.end());
+        flat_vec.push_back('\0'); // Null-terminate each string
+    }
+    auto tensor = torch::tensor(flat_vec, torch::kInt8);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+void PytorchSerDes::setFeature(const std::string &Name, const std::vector<bool> &Value) {
+    std::vector<uint8_t> bool_vec(Value.begin(), Value.end());
+    auto tensor = torch::tensor(bool_vec, torch::kUInt8);
+    this->inputTensors->push_back(tensor.clone());
+}
+
+// void PytorchSerDes::setRequest(void *Request) {
+//     CompiledModel = reinterpret_cast<torch::inductor::AOTIModelContainerRunnerCpu *>(Request);
+// }
+
+void PytorchSerDes::cleanDataStructures() {
+    this->inputTensors->clear(); // Clear the input vector
+}
+
+void *PytorchSerDes::deserializeUntyped(void *Data) {
+    if (Data == nullptr) {
+        return nullptr;
+    }
+
+    // Assume Data is a pointer to a vector of tensors
+    std::vector<torch::Tensor> *serializedTensors = reinterpret_cast<std::vector<torch::Tensor> *>(Data);
+
+    if (serializedTensors->empty()) {
+        return nullptr;
+    }
+
+    auto type_vect = serializedTensors->at(0).dtype();
+
+    if (type_vect == torch::kInt32) {
+        return copyTensorToVect<int32_t>(serializedTensors);
+    } 
+    else if (type_vect == torch::kInt64) {
+        return copyTensorToVect<int64_t>(serializedTensors);
+    } 
+    else if (type_vect == torch::kFloat32) {
+        return copyTensorToVect<float>(serializedTensors);
+    } 
+    else if (type_vect == torch::kFloat64) {
+        return copyTensorToVect<double>(serializedTensors);
+    } 
+    else if (type_vect == torch::kBool) {
+        return copyTensorToVect<bool>(serializedTensors);
+    } 
+    else if (type_vect == torch::kInt8) {
+        return copyTensorToVect<int8_t>(serializedTensors);
+    } 
+    else {
+        llvm::errs() << "Unsupported tensor dtype.\n";
+        return nullptr;
+    }
+}
+
+void *PytorchSerDes::getSerializedData() {
+    std::vector<torch::Tensor> serializedData = *(this->outputTensors);
+
+    // Allocate memory for the output and copy the serialized data
+    auto *output = new std::vector<torch::Tensor>(serializedData);
+    return static_cast<void *>(output);
+}
+
+
+
+} // namespace MLBridge
diff --git a/include/MLModelRunner/MLModelRunner.h b/include/MLModelRunner/MLModelRunner.h
index 77892ec..a764193 100644
--- a/include/MLModelRunner/MLModelRunner.h
+++ b/include/MLModelRunner/MLModelRunner.h
@@ -49,6 +49,7 @@
 #ifndef C_LIBRARY
 #include "SerDes/protobufSerDes.h"
 #include "SerDes/tensorflowSerDes.h"
+#include "SerDes/pytorchSerDes.h"
 #endif
 namespace MLBridge {
 
@@ -81,7 +82,7 @@ class MLModelRunner {
   }
 
   /// Type of the MLModelRunner
-  enum class Kind : int { Unknown, Pipe, gRPC, ONNX, TFAOT };
+  enum class Kind : int { Unknown, Pipe, gRPC, ONNX, TFAOT, PTAOT};
 
   Kind getKind() const { return Type; }
   SerDesKind getSerDesKind() const { return SerDesType; }
@@ -151,11 +152,14 @@ class MLModelRunner {
       SerDes = std::make_unique<BitstreamSerDes>();
       break;
 #ifndef C_LIBRARY
-    case SerDesKind::Protobuf:
-      SerDes = std::make_unique<ProtobufSerDes>();
-      break;
-    case SerDesKind::Tensorflow:
-      SerDes = std::make_unique<TensorflowSerDes>();
+    // case SerDesKind::Protobuf:
+    //   SerDes = std::make_unique<ProtobufSerDes>();
+    //   break;
+    // case SerDesKind::Tensorflow:
+    //   SerDes = std::make_unique<TensorflowSerDes>();
+    //   break;
+    case SerDesKind::Pytorch:
+      SerDes = std::make_unique<PytorchSerDes>();
       break;
 #endif
     case SerDesKind::Unknown:
diff --git a/include/MLModelRunner/ONNXModelRunner/ONNXModelRunner.h b/include/MLModelRunner/ONNXModelRunner/ONNXModelRunner.h
old mode 100755
new mode 100644
diff --git a/include/MLModelRunner/ONNXModelRunner/utils.h b/include/MLModelRunner/ONNXModelRunner/utils.h
old mode 100755
new mode 100644
diff --git a/include/MLModelRunner/PTModelRunner.h b/include/MLModelRunner/PTModelRunner.h
new file mode 100644
index 0000000..09eaf1a
--- /dev/null
+++ b/include/MLModelRunner/PTModelRunner.h
@@ -0,0 +1,89 @@
+#ifndef PTMODELRUNNER_H
+#define PTMODELRUNNER_H
+
+#include "MLModelRunner/MLModelRunner.h"
+#include "SerDes/TensorSpec.h"
+#include "SerDes/pytorchSerDes.h"
+#include "llvm/Support/ErrorHandling.h"
+#include <torch/torch.h>
+#include <torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h> // or model_container_runner_cuda.h for CUDA
+
+#include <memory>
+#include <vector>
+
+namespace MLBridge
+{
+
+  /// PTModelRunner - PyTorch Compiled model implementation of the
+  /// MLModelRunner. It uses an AOT-compiled model for efficient execution.
+  class PTModelRunner final : public MLModelRunner
+  {
+  public:
+    // New constructor that takes the model path as an input
+    PTModelRunner(const std::string &modelPath, llvm::LLVMContext &Ctx)
+        : MLModelRunner(MLModelRunner::Kind::PTAOT, SerDesKind::Pytorch, &Ctx)
+    {
+      this->SerDes = new PytorchSerDes();
+
+      c10::InferenceMode mode;
+      this->CompiledModel = new torch::inductor::AOTIModelContainerRunnerCpu(modelPath);
+    }
+
+    virtual ~PTModelRunner() = default;
+
+    virtual void requestExit() override
+    {
+      llvm_unreachable("requestExit() is not supported in PTModelRunner");
+    }
+
+    static bool classof(const MLModelRunner *R)
+    {
+      return R->getKind() == MLModelRunner::Kind::PTAOT;
+    }
+
+    template <typename U, typename T, typename... Types>
+    void populateFeatures(const std::pair<U, T> &var1,
+                          const std::pair<U, Types> &...var2);
+
+    void populateFeatures() {}
+
+    void *evaluateUntyped() override;
+    PytorchSerDes *SerDes;
+    // Compiled model container added to the PTModelRunner
+    torch::inductor::AOTIModelContainerRunnerCpu *CompiledModel; 
+  };
+
+  void *PTModelRunner::evaluateUntyped()
+  {
+
+    if ((*(this->SerDes->inputTensors)).empty())
+    {
+      llvm::errs() << "Input vector is empty.\n";
+      return nullptr;
+    }
+
+    try
+    {
+      auto outputs = this->CompiledModel->run((*(this->SerDes->inputTensors)));
+      for (auto i = outputs.begin(); i != outputs.end(); ++i)
+        (*(this->SerDes->outputTensors)).push_back(*i);
+      void *rawData = this->SerDes->deserializeUntyped(this->SerDes->outputTensors);
+      return rawData;
+    }
+    catch (const c10::Error &e)
+    {
+      llvm::errs() << "Error during model evaluation: " << e.what() << "\n";
+      return nullptr;
+    }
+  }
+
+  template <typename U, typename T, typename... Types>
+  void PTModelRunner::populateFeatures(const std::pair<U, T> &var1,
+                                       const std::pair<U, Types> &...var2)
+  {
+    SerDes->setFeature(var1.first, var1.second);
+    PTModelRunner::populateFeatures(var2...);
+  }
+} // namespace MLBridge
+
+#endif // TFMODELRUNNER_H
diff --git a/include/MLModelRunner/PipeModelRunner.h b/include/MLModelRunner/PipeModelRunner.h
old mode 100755
new mode 100644
diff --git a/include/SerDes/baseSerDes.h b/include/SerDes/baseSerDes.h
index 977143a..6966f64 100644
--- a/include/SerDes/baseSerDes.h
+++ b/include/SerDes/baseSerDes.h
@@ -43,7 +43,7 @@ namespace MLBridge {
 /// communication by the MLModelRunner.
 /// Currently, (int, float) or (long, double), char and bool are supported.
 /// Vectors of these types are supported as well.
-enum class SerDesKind : int { Unknown, Json, Bitstream, Protobuf, Tensorflow };
+enum class SerDesKind : int { Unknown, Json, Bitstream, Protobuf, Tensorflow, Pytorch };
 class BaseSerDes {
 public:
   SerDesKind getKind() const { return Type; }
diff --git a/include/SerDes/protobufSerDes.h b/include/SerDes/protobufSerDes.h
index 25b7e3f..d7a5f3a 100644
--- a/include/SerDes/protobufSerDes.h
+++ b/include/SerDes/protobufSerDes.h
@@ -18,7 +18,7 @@
 #include "google/protobuf/extension_set.h"
 #include "google/protobuf/message.h"
 
-using namespace google::protobuf;
+// using namespace google::protobuf;
 
 namespace MLBridge {
 /// ProtobufSerDes - Protobuf Serialization/Deserialization to support gRPC
@@ -54,12 +54,12 @@ class ProtobufSerDes : public BaseSerDes {
   void *getSerializedData() override;
   void cleanDataStructures() override;
 
-  Message *getMessage() { return Response; };
+  google::protobuf::Message *getMessage() { return Response; };
 
 private:
   void *deserializeUntyped(void *data) override;
-  Message *Response;
-  Message *Request;
+  google::protobuf::Message *Response;
+  google::protobuf::Message *Request;
 };
 } // namespace MLBridge
 
diff --git a/include/SerDes/pytorchSerDes.h b/include/SerDes/pytorchSerDes.h
new file mode 100644
index 0000000..66845ec
--- /dev/null
+++ b/include/SerDes/pytorchSerDes.h
@@ -0,0 +1,56 @@
+#ifndef PYTORCH_SERIALIZER_H
+#define PYTORCH_SERIALIZER_H
+
+#include "SerDes/baseSerDes.h"
+#include <iostream>
+#include <vector>
+
+// #include <torch/script.h>
+#include <torch/torch.h>
+#include <torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h>
+
+namespace MLBridge {
+class PytorchSerDes : public BaseSerDes {
+public:
+  // Constructor that takes the model path and loads the model
+  PytorchSerDes() : BaseSerDes(SerDesKind::Pytorch) {
+    inputTensors = std::make_shared<std::vector<torch::Tensor>>();
+    outputTensors = new std::vector<torch::Tensor>();
+  }
+
+  ~PytorchSerDes() {}
+
+  static bool classof(const BaseSerDes *S) {
+    return S->getKind() == SerDesKind::Pytorch;
+  }
+
+#define SET_FEATURE(TYPE, _)                                                      \
+  void setFeature(const std::string &, const TYPE) override;                 \
+  void setFeature(const std::string &, const std::vector<TYPE> &) override;
+  SUPPORTED_TYPES(SET_FEATURE)
+#undef SET_FEATURE
+
+  void setRequest(void *request){
+  }
+
+ template <typename T>
+    std::vector<T> *copyTensorToVect(std::vector<torch::Tensor> *serializedTensors) {
+        auto *ret = new std::vector<T>();
+        for (const auto &tensor : *serializedTensors) {
+            ret->insert(ret->end(), tensor.data_ptr<T>(), tensor.data_ptr<T>() + tensor.numel());
+        }
+        return ret;
+    }
+  
+
+  void *getSerializedData() override;
+  void cleanDataStructures() override;
+  void *deserializeUntyped(void *data) override;
+
+  std::vector<torch::Tensor> *outputTensors; // Storage for the PyTorch output tensor
+  std::shared_ptr<std::vector<torch::Tensor>> inputTensors; // Declaration of the input tensor
+
+};
+} // namespace MLBridge
+
+#endif
\ No newline at end of file
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index f4a83a3..c37107e 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -1,11 +1,17 @@
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-rtti")
+# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-rtti")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ")
+
+find_package(Torch REQUIRED)
 
 add_library(MLBridgeCPPTest OBJECT MLBridgeTest.cpp)
-file(GLOB MODEL_OBJECTS ${CMAKE_CURRENT_SOURCE_DIR}/tf_models/*.o)
 
-foreach(MODEL_OBJECT ${MODEL_OBJECTS})
-  target_link_libraries(MLBridgeCPPTest PRIVATE ${MODEL_OBJECT})
-endforeach()
+# file(GLOB MODEL_OBJECTS ${CMAKE_CURRENT_SOURCE_DIR}/tf_models/*.o)
+
+# foreach(MODEL_OBJECT ${MODEL_OBJECTS})
+#   target_link_libraries(MLBridgeCPPTest PRIVATE ${MODEL_OBJECT})
+# endforeach()
 target_link_libraries(MLBridgeCPPTest PRIVATE MLCompilerBridge )
 target_include_directories(MLBridgeCPPTest PRIVATE ${CMAKE_BINARY_DIR}/include ${TENSORFLOW_AOT_PATH}/include ${CMAKE_CURRENT_SOURCE_DIR}/include)
-target_link_libraries(MLBridgeCPPTest PRIVATE tf_xla_runtime)
+target_link_libraries(MLBridgeCPPTest PRIVATE tf_xla_runtime "${TORCH_LIBRARIES}")
+
+set_property(TARGET MLBridgeCPPTest PROPERTY CXX_STANDARD 17)
diff --git a/test/MLBridgeTest.cpp b/test/MLBridgeTest.cpp
index f879de9..e8c5108 100644
--- a/test/MLBridgeTest.cpp
+++ b/test/MLBridgeTest.cpp
@@ -6,283 +6,99 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "HelloMLBridge_Env.h"
 #include "MLModelRunner/MLModelRunner.h"
-#include "MLModelRunner/ONNXModelRunner/ONNXModelRunner.h"
-#include "MLModelRunner/PipeModelRunner.h"
-#include "MLModelRunner/TFModelRunner.h"
-#include "MLModelRunner/Utils/DataTypes.h"
-#include "MLModelRunner/Utils/MLConfig.h"
-#include "MLModelRunner/gRPCModelRunner.h"
-#include "ProtosInclude.h"
+#include "MLModelRunner/PTModelRunner.h"
+#include "SerDes/pytorchSerDes.h"
 #include "llvm/Support/CommandLine.h"
 #include <algorithm>
+#include <chrono>
 #include <fstream>
 #include <iostream>
 #include <iterator>
 #include <memory>
 #include <random>
 #include <string>
+#include <vector>
 
-#define debug_out                                                              \
-  if (!silent)                                                                 \
-  std::cout
-using namespace grpc;
+using namespace MLBridge;
+using namespace std;
 
-static llvm::cl::opt<std::string>
-    cl_server_address("test-server-address", llvm::cl::Hidden,
-                      llvm::cl::desc("Server address, format <ip>:<port>"),
-                      llvm::cl::init(""));
-
-static llvm::cl::opt<std::string>
-    cl_pipe_name("test-pipe-name", llvm::cl::Hidden, llvm::cl::init(""),
-                 llvm::cl::desc("Name for pipe file"));
-
-static llvm::cl::opt<std::string>
-    cl_onnx_path("onnx-model-path", llvm::cl::Hidden, llvm::cl::init(""),
-                 llvm::cl::desc("Path to onnx model"));
-
-static llvm::cl::opt<std::string> cl_test_config(
-    "test-config", llvm::cl::Hidden,
-    llvm::cl::desc("Method for communication with python model"));
-
-static llvm::cl::opt<bool>
-    silent("silent", llvm::cl::Hidden, llvm::cl::init(false),
-           llvm::cl::desc("Only print errors when if set to true"));
-
-namespace {
-std::string basename;
-SerDesKind SerDesType;
-
-std::string test_config;
-std::string pipe_name;
-std::string server_address;
-std::string onnx_path;
-
-// send value of type T1. Test received value of type T2 against expected value
-template <typename T1, typename T2, typename MLRunnerTy>
-void testPrimitive(MLRunnerTy &MLRunner, std::string label, T1 value,
-                   T2 expected) {
-  std::pair<std::string, T1> p("request_" + label, value);
-  MLRunner->populateFeatures(p);
-  T2 out = MLRunner->template evaluate<T2>();
-  debug_out << "  " << label << " reply: " << out << "\n";
-  if (std::abs(out - expected) > 10e-6) {
-    std::cerr << "Error: Expected " << label << " reply: " << expected
-              << ", Received: " << out << "\n";
-    exit(1);
-  }
-}
-
-template <typename T1, typename T2, typename MLRunnerTy>
-void testVector(MLRunnerTy &MLRunner, std::string label, std::vector<T1> value,
-                std::vector<T2> expected) {
-  std::pair<std::string, std::vector<T1>> p("request_" + label, value);
-  MLRunner->populateFeatures(p);
-  T2 *out;
-  size_t size;
-  MLRunner->template evaluate<T2 *>(out, size);
-  std::vector<T2> reply(out, out + size);
-  debug_out << "  " << label << " reply: ";
-  int i = 0;
-  for (auto x : reply) {
-    debug_out << x << " ";
-    if (std::abs(x - expected[i]) > 10e-6) {
-      std::cerr << "Error: Expected " << label << " reply: " << expected[i]
-                << ", Received: " << x << "\n";
-      exit(1);
-    }
-    i++;
-  }
-  debug_out << "\n";
-}
-
-int testPipeBytes() {
-  if (pipe_name == "") {
-    std::cerr
-        << "Pipe name must be specified via --test-pipe-name=<filename>\n";
-    exit(1);
-  }
-  basename = "./" + pipe_name;
-  SerDesType = SerDesKind::Bitstream;
-  auto MLRunner = std::make_unique<PipeModelRunner>(
-      basename + ".out", basename + ".in", SerDesType, nullptr);
-  testPrimitive(MLRunner, "int", 11, 12);
-  testPrimitive(MLRunner, "long", 1234567890l, 1234567891l);
-  testPrimitive(MLRunner, "float", 3.14f, 4.14f);
-  testPrimitive(MLRunner, "double", 0.123456789123456789, 1.123456789123456789);
-  testPrimitive(MLRunner, "char", 'a', 'b');
-  testPrimitive(MLRunner, "bool", true, false);
-  testVector(MLRunner, "vec_int", std::vector<int>{11, 22, 33},
-             std::vector<int>{12, 23, 34});
-  testVector(MLRunner, "vec_long", std::vector<long>{123456780, 222, 333},
-             std::vector<long>{123456780, 123456781, 123456782});
-  testVector(MLRunner, "vec_float", std::vector<float>{11.1, 22.2, 33.3},
-             std::vector<float>{1.11, 2.22, -3.33, 0});
-  testVector(MLRunner, "vec_double",
-             std::vector<double>{-1.1111111111, -2.2222222222, -3.3333333333},
-             std::vector<double>{1.12345678912345670, -1.12345678912345671});
-  return 0;
-}
-
-int testPipeJSON() {
-  if (pipe_name == "") {
-    std::cerr
-        << "Pipe name must be specified via --test-pipe-name=<filename>\n";
-    exit(1);
+int main(int argc, char **argv)
+{
+  if (argc != 2)
+  {
+    cerr << "Usage: " << argv[0] << " <size>" << endl;
+    return 1;
   }
-  basename = "./" + pipe_name;
-  SerDesType = SerDesKind::Json;
-  auto MLRunner = std::make_unique<PipeModelRunner>(
-      basename + ".out", basename + ".in", SerDesType, nullptr);
-  testPrimitive(MLRunner, "int", 11, 12);
-  testPrimitive(MLRunner, "long", 1234567890l, 12345l);
-  testPrimitive(MLRunner, "float", 3.14f, 4.14f);
 
-  // FIXME: doesn't work if expected value is double
-  testPrimitive(MLRunner, "double", 0.123456789123456789,
-                1.123456789123456789f);
-  testPrimitive(MLRunner, "char", 'a', 'b');
-  testPrimitive(MLRunner, "bool", true, false);
-  testVector(MLRunner, "vec_int", std::vector<int>{11, 22, 33},
-             std::vector<int>{12, 23, 34});
+  int sz = stoi(argv[1]);
+  cout << "Running inference for size: " << sz << endl;
 
-  // FIXME: doesn't work if expected value is long
-  testVector(MLRunner, "vec_long", std::vector<long>{123456780, 222, 333},
-             std::vector<int>{6780, 6781, 6782});
-  testVector(MLRunner, "vec_float", std::vector<float>{11.1, 22.2, 33.3},
-             std::vector<float>{1.11, 2.22, -3.33, 0});
+  // Create random input data of size sz
+  vector<float> inputData(sz);
+  random_device rd;
+  mt19937 gen(rd());
+  uniform_real_distribution<float> dis(0.0, 1.0);
+  generate(inputData.begin(), inputData.end(), [&]()
+           { return dis(gen); });
 
-  // FIXME: doesn't work if expected value is double
-  testVector(MLRunner, "vec_double",
-             std::vector<double>{-1.1111111111, -2.2222222222, -3.3333333333},
-             std::vector<float>{1.12345678912345670, -1.12345678912345671});
-  return 0;
-}
-
-void increment_port(int delta) {
-  int split = server_address.find(":");
-  int port = stoi(server_address.substr(split + 1));
-  server_address =
-      server_address.substr(0, split) + ":" + to_string(port + delta);
-}
-
-int testGRPC() {
-#define gRPCModelRunnerInit(datatype)                                          \
-  increment_port(1);                                                           \
-  MLBridgeTestgRPC_##datatype::Reply response;                                 \
-  MLBridgeTestgRPC_##datatype::Request request;                                \
-  auto MLRunner = std::make_unique<                                            \
-      gRPCModelRunner<MLBridgeTestgRPC_##datatype::MLBridgeTestService,        \
-                      MLBridgeTestgRPC_##datatype::MLBridgeTestService::Stub,  \
-                      MLBridgeTestgRPC_##datatype::Request,                    \
-                      MLBridgeTestgRPC_##datatype::Reply>>(                    \
-      server_address, &request, &response, nullptr);                           \
-  MLRunner->setRequest(&request);                                              \
-  MLRunner->setResponse(&response)
+  llvm::LLVMContext Ctx;
+  string modelPath = "models_pt/model_" + to_string(sz) + ".so";
 
-  if (server_address == "") {
-    std::cerr << "Server Address must be specified via "
-                 "--test-server-address=\"<ip>:<port>\"\n";
-    exit(1);
-  }
+  try
   {
-    gRPCModelRunnerInit(int);
-    testPrimitive(MLRunner, "int", 11, 12);
-  }
-  {
-    gRPCModelRunnerInit(long);
-    testPrimitive(MLRunner, "long", 1234567890l, 1234567891l);
-  }
-  {
-    gRPCModelRunnerInit(float);
-    testPrimitive(MLRunner, "float", 3.14f, 4.14f);
-  }
-  {
-    gRPCModelRunnerInit(double);
-    testPrimitive(MLRunner, "double", 0.123456789123456789,
-                  1.123456789123456789);
-  }
-  increment_port(1);
-  {
-    gRPCModelRunnerInit(bool);
-    testPrimitive(MLRunner, "bool", true, false);
-  }
-  {
-    gRPCModelRunnerInit(vec_int);
-    testVector(MLRunner, "vec_int", std::vector<int>{11, 22, 33},
-               std::vector<int>{12, 23, 34});
-  }
-  {
-    gRPCModelRunnerInit(vec_long);
-    testVector(MLRunner, "vec_long", std::vector<long>{123456780, 222, 333},
-               std::vector<long>{123456780, 123456781, 123456782});
-  }
-  {
-    gRPCModelRunnerInit(vec_float);
-    testVector(MLRunner, "vec_float", std::vector<float>{11.1, 22.2, 33.3},
-               std::vector<float>{1.11, 2.22, -3.33, 0});
-  }
-  {
-    gRPCModelRunnerInit(vec_double);
-    testVector(MLRunner, "vec_double",
-               std::vector<double>{-1.1111111111, -2.2222222222, -3.3333333333},
-               std::vector<double>{1.12345678912345670, -1.12345678912345671});
-  }
-#undef gRPCModelRunnerInit
-  return 0;
-}
-
-class ONNXTest : public MLBridgeTestEnv {
-public:
-  int run(int expectedAction) {
-    onnx_path = cl_onnx_path.getValue();
-    if (onnx_path == "") {
-      std::cerr << "ONNX model path must be specified via "
-                   "--onnx-model-path=<filepath>\n";
-      exit(1);
+    // Start measuring time just before loading the model
+    auto start = chrono::high_resolution_clock::now();
+
+    unique_ptr<PTModelRunner> runner = make_unique<PTModelRunner>(modelPath, Ctx);
+
+    // Prepare the input
+    basic_string<char> input_str = "input";
+    pair<string, vector<float> &> inputPair = make_pair(input_str, ref(inputData));
+    runner->populateFeatures(inputPair);
+
+    // Perform inference
+    void *result = runner->evaluateUntyped();
+
+    auto end = chrono::high_resolution_clock::now();
+    chrono::duration<double, milli> inferenceTime = end - start;
+
+    // Check result
+    if (result)
+    {
+      vector<float> *output = reinterpret_cast<vector<float> *>(result);
+      cout << "Output from the model: ";
+      for (auto &v : *output)
+      {
+        cout << v << " ";
+      }
+      cout << endl;
     }
-    FeatureVector.clear();
-    int n = 100;
-    for (int i = 0; i < n; i++) {
-      float delta = (float)(i - expectedAction) / n;
-      FeatureVector.push_back(delta * delta);
+    else
+    {
+      cerr << "Model evaluation failed." << endl;
+      return 1;
     }
 
-    Agent *agent = new Agent(onnx_path);
-    std::map<std::string, Agent *> agents;
-    agents["agent"] = agent;
-    auto MLRunner = std::make_unique<ONNXModelRunner>(this, agents, nullptr);
-    MLRunner->evaluate<int>();
-    if (lastAction != expectedAction) {
-      std::cerr << "Error: Expected action: " << expectedAction
-                << ", Computed action: " << lastAction << "\n";
-      exit(1);
-    }
-    return 0;
+    // // Log the results
+    // ofstream resultsFile("results.csv", ios::app);
+    // if (resultsFile.is_open())
+    // {
+    //   resultsFile << sz << "," << inferenceTime.count() << endl;
+    //   resultsFile.close();
+    // }
+    // else
+    // {
+    //   cerr << "Failed to open results.csv" << endl;
+    // }
+
+    cout << "Inference completed in: " << inferenceTime.count() << " ms" << endl;
+  }
+  catch (const exception &ex)
+  {
+    cerr << "Error: " << ex.what() << endl;
+    return 1;
   }
-};
-
-} // namespace
-
-int main(int argc, char **argv) {
-  llvm::cl::ParseCommandLineOptions(argc, argv);
-  test_config = cl_test_config.getValue();
 
-  if (test_config == "pipe-bytes") {
-    pipe_name = cl_pipe_name.getValue();
-    testPipeBytes();
-  } else if (test_config == "pipe-json") {
-    pipe_name = cl_pipe_name.getValue();
-    testPipeJSON();
-  } else if (test_config == "grpc") {
-    server_address = cl_server_address.getValue();
-    testGRPC();
-  } else if (test_config == "onnx") {
-    ONNXTest t;
-    t.run(20);
-  } else
-    std::cerr << "--test-config must be provided from [pipe-bytes, pipe-json, "
-                 "grpc, onnx]\n";
   return 0;
 }