Unverified Commit 9ad28f63 authored by aoshen524's avatar aoshen524 Committed by GitHub
Browse files

fix(srt): check if sample_indices is not None before usage. (#5633)

parent d7b1ce65
......@@ -335,13 +335,13 @@ class LogitsProcessor(nn.Module):
aux_pruned_states = torch.cat(aux_pruned_states, dim=-1)
hidden_states_to_store = (
aux_pruned_states[sample_indices]
if sample_indices
if sample_indices is not None
else aux_pruned_states
)
else:
hidden_states_to_store = (
pruned_states[sample_indices]
if sample_indices
if sample_indices is not None
else pruned_states
)
else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment