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

fix bugs in sampler

Browse Source
This commit is contained in:
yuehuayingxueluo
2024-01-04 15:03:18 +08:00
committed by FrankLeeeee
parent 02c1bf8b2a
commit bbfebfb9fc
4 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/inference/sampler.py
View File

@@ -21,7 +21,7 @@ def multinomial_sample(
"""
Sample tokens in a random phase.
"""
random_results = torch.multinomial(probs, num_samples=1, replacement=True).cpu()
random_results = torch.multinomial(probs, num_samples=1).squeeze(1)
return random_results