fix(server): Check for device type correctly when determining initial padding (IBM#16)

njhill · web-flow · commit 686cc6671705 · 2022-12-30T19:30:42.000+01:00
AFAIK there is no torch device type called "gpu".
diff --git a/server/text_generation/models/causal_lm.py b/server/text_generation/models/causal_lm.py
@@ -65,7 +65,7 @@ def from_pb(
             )
             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
@@ -77,7 +77,7 @@ def from_pb(
             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",

Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ def from_pb(`
`65`	`65`	`)`
`66`	`66`	`all_logprobs.append(None)`
`67`	`67`
`68`		`- pad_to_multiple_of = 8 if "gpu" in str(device) else None`
	`68`	`+ pad_to_multiple_of = 8 if device.type == "cuda" else None`
`69`	`69`	`tokenized_inputs = tokenizer(`
`70`	`70`	`inputs,`
`71`	`71`	`return_tensors="pt",`