remove torch inference_mode

HuanzhiMao · Jul 20, 2024 · 5ff790e · 5ff790e
1 parent b84fa39
commit 5ff790e
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/berkeley-function-call-leaderboard/model_handler/oss_handler.py b/berkeley-function-call-leaderboard/model_handler/oss_handler.py
@@ -3,7 +3,6 @@
 
 import ray
 import shortuuid
-import torch
 from eval_checker.eval_checker_constant import FILENAME_INDEX_MAPPING
 from model_handler.handler import BaseHandler
 from model_handler.model_style import ModelStyle
@@ -35,7 +34,6 @@ def _format_prompt(prompt, function, test_category):
         return f"SYSTEM: {SYSTEM_PROMPT}\n{functions}\nUSER: {prompt}\nASSISTANT: "
 
     @ray.remote(num_gpus=1)
-    @torch.inference_mode()
     def _batch_generate(
         question_jsons,
         test_category,