Unverified Commit b4cf832c authored by Nick Hill's avatar Nick Hill Committed by GitHub
Browse files

fix(server): fix reshaping of bloom past_key_values in concatenate() (#252)

Introduced in #214 

Fixes #249
parent db2b4e07
......@@ -335,7 +335,7 @@ class CausalLMBatch(Batch):
[t.view(len(batch), -1, *t.shape[-2:]) for t in layer]
for layer in batch.past_key_values
]
elif batch.past_key_values[0][0].shape == 3:
elif len(batch.past_key_values[0][0].shape) == 3:
for layer in batch.past_key_values:
for k, t in enumerate(layer):
layer[k] = t.view(len(batch), -1, *t.shape[-2:])
......
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