del to pop

InternLM · Mar 5, 2024 · 9f1ce63 · 9f1ce63
1 parent 8476f5b
commit 9f1ce63
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py
@@ -673,13 +673,13 @@ async def __long_context_forward(inputs):
                 if token_count == 0 and slen > max_prefill_token_num:
                     tmp_out = await __long_context_single_forward(inputs, idx)
                     logits_gather.gather(tmp_out)
-                    del tmp_out
+                    tmp_out.pop('logits', None)
                     idx += 1
                 elif token_count + slen > max_prefill_token_num:
                     tmp_out = await __long_context_batched_forward(
                         inputs, indices[0], idx)
                     logits_gather.gather(tmp_out)
-                    del tmp_out
+                    tmp_out.pop('logits', None)
                     indices = []
                     token_count = 0
                 else: