move tensors to cpu for json conversion

Signed-off-by: jagadeesh <jagadeeshj@ideas2it.com>
pytorch · Aug 25, 2023 · 46fadcd · 46fadcd
1 parent beefa50
commit 46fadcd
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/cpp/build.sh b/cpp/build.sh
@@ -25,7 +25,6 @@ function install_dependencies_linux() {
     autoconf \
     automake \
     git \
-    cmake \
     m4 \
     g++ \
     flex \
@@ -175,6 +174,14 @@ function install_libtorch() {
         wget https://download.pytorch.org/libtorch/cu116/libtorch-cxx11-abi-shared-with-deps-1.12.1%2Bcu116.zip
         unzip libtorch-cxx11-abi-shared-with-deps-1.12.1+cu116.zip
         rm libtorch-cxx11-abi-shared-with-deps-1.12.1+cu116.zip
+      elif [ "$CUDA" = "cu117" ]; then
+        wget https://download.pytorch.org/libtorch/cu117/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcu117.zip
+        unzip libtorch-cxx11-abi-shared-with-deps-2.0.1+cu117.zip
+        rm libtorch-cxx11-abi-shared-with-deps-2.0.1+cu117.zip
+      elif [ "$CUDA" = "cu118" ]; then
+        wget https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcu118.zip
+        unzip libtorch-cxx11-abi-shared-with-deps-2.0.1+cu118.zip
+        rm libtorch-cxx11-abi-shared-with-deps-2.0.1+cu118.zip
       else
         wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.12.1%2Bcpu.zip
         unzip libtorch-cxx11-abi-shared-with-deps-1.12.1+cpu.zip
@@ -246,7 +253,7 @@ function build() {
   if [ "$CUDA" != "" ]; then
     MAYBE_CUDA_COMPILER='-DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc'
   fi
-  
+
   # Build torchserve_cpp with cmake
   cd "$BWD" || exit
   YAML_CPP_CMAKE_DIR=$DEPS_DIR/yaml-cpp-build
@@ -265,7 +272,7 @@ function build() {
     "$MAYBE_CUDA_COMPILER"                                                                    \
     ..
 
-    if [ "$CUDA" = "cu102" ] || [ "$CUDA" = "cu113" ] || [ "$CUDA" = "cu116" ]; then
+    if [ "$CUDA" = "cu102" ] || [ "$CUDA" = "cu113" ] || [ "$CUDA" = "cu116" ] || [ "$CUDA" = "cu117" ] || [ "$CUDA" = "cu118" ]; then
       export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/bin/nvcc
     fi
   elif [ "$PLATFORM" = "Mac" ]; then
@@ -333,7 +340,7 @@ INSTALL_DEPENDENCIES=false
 PREFIX=""
 COMPILER_FLAGS=""
 CUDA=""
-USAGE="./build.sh [-j num_jobs] [-g cu102|cu113|cu116] [-q|--with-quic] [--install-dependencies] [-p|--prefix] [-x|--compiler-flags]"
+USAGE="./build.sh [-j num_jobs] [-g cu102|cu113|cu116|cu117|cu118] [-q|--with-quic] [--install-dependencies] [-p|--prefix] [-x|--compiler-flags]"
 while [ "$1" != "" ]; do
   case $1 in
     -j | --jobs ) shift
@@ -395,7 +402,7 @@ cd $BASE_DIR
 
 install_folly
 install_kineto
-#install_libtorch
+install_libtorch
 install_yaml_cpp
 build
 symlink_torch_libs

diff --git a/cpp/src/examples/image_classifier/resnet-18/resnet-18_handler.cc b/cpp/src/examples/image_classifier/resnet-18/resnet-18_handler.cc
@@ -176,13 +176,16 @@ void ResnetHandler::Postprocess(
       torch::Tensor ps = F::softmax(data, F::SoftmaxFuncOptions(1));
       std::tuple<torch::Tensor, torch::Tensor> result =
           torch::topk(ps, kTopKClasses, 1, true, true);
-      auto [probs, classes] = result;
+      torch::Tensor probs = std::get<0>(result);
+      torch::Tensor classes = std::get<1>(result);
 
+      probs = probs.to(torch::kCPU);
+      classes = classes.to(torch::kCPU);
       // Convert tensors to C++ vectors
-      std::vector<float> probs_vector(probs.data<float>(),
-                                      probs.data<float>() + probs.numel());
-      std::vector<long> classes_vector(classes.data<long>(),
-                                       classes.data<long>() + classes.numel());
+      std::vector<float> probs_vector(probs.data_ptr<float>(),
+                                      probs.data_ptr<float>() + probs.numel());
+      std::vector<long> classes_vector(
+          classes.data_ptr<long>(), classes.data_ptr<long>() + classes.numel());
 
       // Create a JSON object using folly::dynamic
       folly::dynamic json_response = folly::dynamic::object;