PR changes

Signed-off-by: Abhishek <[email protected]>
BernardZach · Oct 11, 2024 · 39d2868 · 39d2868
1 parent d2796f6
commit 39d2868
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/processing_utils.py b/src/transformers/processing_utils.py
@@ -27,7 +27,6 @@
 from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union
 
 import numpy as np
-import torch
 import typing_extensions
 
 from .dynamic_module_utils import custom_object_save
@@ -93,6 +92,7 @@ class FlashAttentionKwargs(TypedDict, total=False):
         max_length_k (`int`, *optional*):
             Maximum sequence length for key state.
     """
+    import torch
 
     cu_seq_lens_q: Optional[torch.LongTensor]
     cu_seq_lens_k: Optional[torch.LongTensor]