Skip to content

Commit 167a154

Browse files
committed
Fix PT compile warmup
Inadvertently broken by new dtype positional arg added to Batch.from_pb()
1 parent 47e8ba7 commit 167a154

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

server/text_generation_server/utils/warmup.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ def __eval_shape(batch_size: int, sequence_length: int, num_new_tokens: int, as_
5252
batch, errors = model.batch_type.from_pb(
5353
request.batch,
5454
tokenizer=model.tokenizer,
55+
dtype=model.dtype,
5556
device=model.device,
5657
embeddings_lookup=model.word_embeddings,
5758
prefix_cache=model.prefix_cache,

0 commit comments

Comments
 (0)