github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

precision alignment

Browse Source
This commit is contained in:
yuehuayingxueluo
2024-01-02 18:30:11 +08:00
committed by FrankLeeeee
parent 62968588d1
commit 9489dc64d8
5 changed files with 45 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/inference/struct.py
View File

@@ -308,7 +308,7 @@ class BatchInfo:
input_len_list.append(1)
return torch.tensor(input_list, dtype=torch.long, device=self.device), torch.tensor(
input_len_list, dtype=torch.int, device=device
input_len_list, dtype=torch.int, device=self.device
)
def get_sequence_lengths(self):