Просмотр исходного кода

model-conversion : run-org-model.py fails to run on mac m1 (#16213)

Signed-off-by: Jie Fu <jiefu@tencent.com>
Jie Fu (傅杰) 3 месяцев назад
Родитель
Сommit
7735706b93
1 измененных файлов с 1 добавлено и 1 удалено
  1. 1 1
      examples/model-conversion/scripts/causal/run-org-model.py

+ 1 - 1
examples/model-conversion/scripts/causal/run-org-model.py

@@ -193,7 +193,7 @@ print(f"Input text: {repr(prompt)}")
 print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}")
 
 with torch.no_grad():
-    outputs = model(input_ids)
+    outputs = model(input_ids.to(model.device))
     logits = outputs.logits
 
     # Extract logits for the last token (next token prediction)