From 722eb39526ec39e4615f3a631e08c487c4c17f1f Mon Sep 17 00:00:00 2001
From: cdpath <11472839+cdpath@users.noreply.github.com>
Date: Tue, 12 Sep 2023 10:41:03 +0800
Subject: [PATCH] fix potential oom issue (#387)

---
 opencompass/openicl/icl_retriever/icl_mdl_retriever.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/opencompass/openicl/icl_retriever/icl_mdl_retriever.py b/opencompass/openicl/icl_retriever/icl_mdl_retriever.py
index 43fe12d1..f92e1acf 100644
--- a/opencompass/openicl/icl_retriever/icl_mdl_retriever.py
+++ b/opencompass/openicl/icl_retriever/icl_mdl_retriever.py
@@ -141,6 +141,7 @@ class MDLRetriever(TopkRetriever):
         """Retrieve the in-context example index for each test example."""
         return self.topk_search()
 
+    @torch.no_grad()
     def cal_ce(self, input_texts: List[str], mask_length=None):
         if self.metric_model is None:
             logger.info(