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): ...@@ -335,13 +335,13 @@ class LogitsProcessor(nn.Module):
aux_pruned_states = torch.cat(aux_pruned_states, dim=-1) aux_pruned_states = torch.cat(aux_pruned_states, dim=-1)
hidden_states_to_store = ( hidden_states_to_store = (
aux_pruned_states[sample_indices] aux_pruned_states[sample_indices]
if sample_indices if sample_indices is not None
else aux_pruned_states else aux_pruned_states
) )
else: else:
hidden_states_to_store = ( hidden_states_to_store = (
pruned_states[sample_indices] pruned_states[sample_indices]
if sample_indices if sample_indices is not None
else pruned_states else pruned_states
) )
else: 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