fix the crashing when there are zero-size inputs (#120)

* fix the crashing when there are zero-size inputs

* Typo

---------

Co-authored-by: Iman Tabrizian <[email protected]>

Author: yongbinfeng

src/libtorch.cc

@@ -2136,10 +2136,18 @@ ModelInstanceState::SetInputTensors(
 
       (*input_tensors)[input_index_map_[input_name]] = input_list;
     } else {
-      // Remove constness to align with the signature of torch::from_blob()
-      torch::Tensor input_tensor = torch::from_blob(
-          const_cast<char*>(input_buffer), batchn_shape, updated_options);
-      (*input_tensors)[input_index_map_[input_name]] = input_tensor;
+      if (batchn_byte_size) {
+        // Remove constness to align with the signature of torch::from_blob()
+        torch::Tensor input_tensor = torch::from_blob(
+            const_cast<char*>(input_buffer), batchn_shape, updated_options);
+        (*input_tensors)[input_index_map_[input_name]] = input_tensor;
+      } else {
+        // torch:from_blob seems not working when the input size is 0
+        // create zero-length inputs directly
+        torch::Tensor input_tensor =
+            torch::zeros(batchn_shape, updated_options);
+        (*input_tensors)[input_index_map_[input_name]] = input_tensor;
+      }
     }
   }
 
@@ -2168,9 +2176,15 @@ ModelInstanceState::SetInputTensors(
                                  ? options.device(torch::kCUDA, device_.index())
                                  : options.device(torch::kCPU);
 
-      torch::Tensor input_tensor = torch::from_blob(
-          const_cast<char*>(dst_buffer), shape, updated_options);
-      (*input_tensors)[input_index_map_[input_name]] = input_tensor;
+      if (dst_buffer_byte_size) {
+        torch::Tensor input_tensor = torch::from_blob(
+            const_cast<char*>(dst_buffer), shape, updated_options);
+        (*input_tensors)[input_index_map_[input_name]] = input_tensor;
+      } else {
+        // special handle when input has zero size
+        torch::Tensor input_tensor = torch::zeros(shape, updated_options);
+        (*input_tensors)[input_index_map_[input_name]] = input_tensor;
+      }
     }
   }