From 686cc6671705c666b767fffe71b2ed9c9b6fccd1 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 30 Dec 2022 10:30:42 -0800 Subject: [PATCH] fix(server): Check for device type correctly when determining initial padding (#16) AFAIK there is no torch device type called "gpu". --- server/text_generation/models/causal_lm.py | 2 +- server/text_generation/models/seq2seq_lm.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/server/text_generation/models/causal_lm.py b/server/text_generation/models/causal_lm.py index aeecf12..3ad3621 100644 --- a/server/text_generation/models/causal_lm.py +++ b/server/text_generation/models/causal_lm.py @@ -65,7 +65,7 @@ class CausalLMBatch: ) all_logprobs.append(None) - pad_to_multiple_of = 8 if "gpu" in str(device) else None + pad_to_multiple_of = 8 if device.type == "cuda" else None tokenized_inputs = tokenizer( inputs, return_tensors="pt", diff --git a/server/text_generation/models/seq2seq_lm.py b/server/text_generation/models/seq2seq_lm.py index fc80c60..4095db9 100644 --- a/server/text_generation/models/seq2seq_lm.py +++ b/server/text_generation/models/seq2seq_lm.py @@ -77,7 +77,7 @@ class Seq2SeqLMBatch: decoder_logprobs.append(None) # Tokenize batch - pad_to_multiple_of = 8 if "gpu" in str(device) else None + pad_to_multiple_of = 8 if device.type == "cuda" else None tokenized_inputs = tokenizer( inputs, return_tensors="pt",