Setting InferenceMode guard

Signed-off-by: Shrinath Suresh <shrinath@ideas2it.com>
pytorch · Sep 6, 2023 · bd03fd8 · bd03fd8
1 parent 0d12619
commit bd03fd8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/cpp/src/examples/babyllama/baby_llama_handler.cc b/cpp/src/examples/babyllama/baby_llama_handler.cc
@@ -160,6 +160,7 @@ torch::Tensor LlmHandler::Inference(
     std::shared_ptr<torch::Device>& device,
     std::pair<std::string&, std::map<uint8_t, std::string>&>& idx_to_req_id,
     std::shared_ptr<torchserve::InferenceResponseBatch>& response_batch) {
+  c10::InferenceMode guard;
   std::vector<torch::Tensor> batch_output_vector;
   for (const torch::jit::IValue& input : inputs) {
     std::vector<torch::Tensor> tensor_vector;