riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-02 19:21:31 +06:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

Fix passing str dtype to static cache (#33741)

Browse Source 
Co-authored-by: Guang Yang <guangyang@fb.com>
This commit is contained in:
Guang Yang 2024-10-01 00:50:17 -07:00 committed by GitHub
parent c269c5c74d
commit 808997a634
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/integrations/executorch.py
View File

@ -68,7 +68,7 @@ class TorchExportableModuleWithStaticCache(torch.nn.Module):
config=self.model.config,
batch_size=self.model.generation_config.cache_config.batch_size,
max_cache_len=self.model.generation_config.cache_config.max_cache_len,
dtype=self.model.config.torch_dtype,
dtype=self.model.dtype,
)
self.is_causal = any("CausalLM" in arch for arch in self.model.config.architectures)
if self.is_causal:

2
tests/utils/test_cache_utils.py
View File

@ -181,7 +181,7 @@ class CacheTest(unittest.TestCase):
set_seed(0)
device = "cpu"
dtype = torch.float32
dtype = "bfloat16"
cache_implementation = "static"
attn_implementation = "sdpa" # Export and ExecuTorch only works for SdpaAttention
batch_size = 1