fix: correct torch.randint upper bound to include both values (#2733)

This commit is contained in:
Harikrishna KP 2026-04-18 14:49:55 +05:30 committed by GitHub
parent 445d18ccce
commit 938f05fce8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -262,7 +262,7 @@ def make_reject_y(y_o, y_lens):
reject_y = []
reject_y_lens = []
for b in range(bs):
process_item_idx = torch.randint(0, 1, size=(1,))[0]
process_item_idx = torch.randint(0, 2, size=(1,))[0]
if process_item_idx == 0:
new_y = repeat_P(y_o[b])
reject_y.append(new_y)