huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

cast to int32

Browse Source
This commit is contained in:
OlivierDehaene 2024-10-24 19:01:40 +02:00
parent ea66379e3c
commit d1e95ceaff
No known key found for this signature in database
GPG Key ID: BB104D67809DA93C
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/models/flash_causal_lm.py
View File

@ -945,7 +945,7 @@ class FlashCausalLMBatch(Batch):
) )
self.cu_seqlen_prefill = torch.nn.functional.pad( self.cu_seqlen_prefill = torch.nn.functional.pad(
torch.cumsum(self.input_lengths_tensor, dim=0), (1, 0) torch.cumsum(self.input_lengths_tensor, dim=0), (1, 0)
) ).to(torch.int32)
self.cache_lengths_tensor = torch.tensor( self.cache_lengths_tensor = torch.tensor(
self.cache_lengths, dtype=torch.int32, device=device self.cache_lengths, dtype=torch.int32, device=device
) )