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

…252)

Introduced in #214 

Fixes #249

server/text_generation_server/models/causal_lm.py

@@ -335,7 +335,7 @@ def concatenate(cls, batches: List["CausalLMBatch"]) -> "CausalLMBatch":
                     [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:])