minor fix doc

This commit is contained in:
yaozengwei 2022-06-11 22:44:06 +08:00
parent 64c9053e39
commit 7f94b86bb0

View File

@ -122,7 +122,7 @@ def stack_states(
conv_caches = []
for layer in state_list[0][1]:
if batch_size > 1:
# Note: We will stack conv_caches[layer][] later to get attn_caches[layer] # noqa
# Note: We will stack conv_caches[layer][] later to get conv_caches[layer] # noqa
conv_caches.append([layer])
else:
conv_caches.append(layer.unsqueeze(0))