From 3e25d60837674a491042aec919d926f011a3c548 Mon Sep 17 00:00:00 2001
From: Konstantinos Parasyris <koparasy@gmail.com>
Date: Thu, 19 Oct 2023 12:57:34 -0700
Subject: [PATCH] Fix model inference tests (#12)

---
 src/ml/surrogate.hpp  | 6 +++---
 tests/CMakeLists.txt  | 2 ++
 tests/torch_model.cpp | 9 +++++----
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/ml/surrogate.hpp b/src/ml/surrogate.hpp
index 8e1f930c..babb5157 100644
--- a/src/ml/surrogate.hpp
+++ b/src/ml/surrogate.hpp
@@ -145,7 +145,7 @@ PERFFASPECT()
   // -------------------------------------------------------------------------
 PERFFASPECT()
   inline void _evaluate(long num_elements,
-                        long num_in,
+                        size_t num_in,
                         size_t num_out,
                         const TypeInValue** inputs,
                         TypeInValue** outputs)
@@ -196,9 +196,9 @@ PERFFASPECT()
 
 PERFFASPECT()
   inline void evaluate(long num_elements,
-                       long num_in,
+                       size_t num_in,
                        size_t num_out,
-                       TypeInValue** inputs,
+                       const TypeInValue** inputs,
                        TypeInValue** outputs)
   {
     _evaluate(num_elements, num_in, num_out, inputs, outputs);
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index ef5d9e15..1b45e5c2 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -10,6 +10,8 @@ function (ADDTEST binary_name src_file test_name)
   target_link_libraries(${binary_name} PRIVATE AMS umpire MPI::MPI_CXX)
 
   add_test(NAME "${test_name}::HOST" COMMAND ${binary_name} 0 ${ARGN})
+
+  target_compile_definitions(${binary_name} PRIVATE ${AMS_APP_DEFINES})
   if (WITH_CUDA)
     set_target_properties(${binary_name} PROPERTIES CUDA_ARCHITECTURES "${AMS_CUDA_ARCH}")
     set_property(TARGET ${binary_name} PROPERTY CUDA_SEPARABLE_COMPILATION ON)
diff --git a/tests/torch_model.cpp b/tests/torch_model.cpp
index 7fe2cb2b..2be4dfef 100644
--- a/tests/torch_model.cpp
+++ b/tests/torch_model.cpp
@@ -15,7 +15,7 @@
 #include <wf/data_handler.hpp>
 #include <wf/resource_manager.hpp>
 
-#define SIZE (32 * 1024 + 3)
+#define SIZE (32L * 1024L + 3L)
 
 template <typename T>
 void inference(char *path, int device, AMSResourceType resource)
@@ -23,7 +23,7 @@ void inference(char *path, int device, AMSResourceType resource)
   using namespace ams;
   SurrogateModel<T> model(path, !device);
 
-  std::vector<T *> inputs;
+  std::vector<const T *> inputs;
   std::vector<T *> outputs;
 
   for (int i = 0; i < 2; i++)
@@ -35,8 +35,9 @@ void inference(char *path, int device, AMSResourceType resource)
   model.evaluate(
       SIZE, inputs.size(), outputs.size(), inputs.data(), outputs.data());
 
+
   for (int i = 0; i < 2; i++)
-    ResourceManager::deallocate(inputs[i], resource);
+    ResourceManager::deallocate(const_cast<T*>(inputs[i]), resource);
 
   for (int i = 0; i < 4; i++)
     ResourceManager::deallocate(outputs[i], resource);
@@ -56,7 +57,7 @@ int main(int argc, char *argv[])
   if (use_device == 1) {
     AMSSetupAllocator(AMSResourceType::DEVICE);
     AMSSetDefaultAllocator(AMSResourceType::DEVICE);
-    AMSResourceType resource = AMSResourceType::DEVICE;
+    resource = AMSResourceType::DEVICE;
   }
 
   inference<double>(model_path, use_device, resource);