We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 616c813 commit a949d97Copy full SHA for a949d97
examples/pytorch/nlp/huggingface_models/language-modeling/quantization/mix-precision/quantize.py
@@ -103,6 +103,9 @@ def initialize_model_and_tokenizer(model_name_or_path):
103
104
model, tokenizer = initialize_model_and_tokenizer(args.model_name_or_path)
105
device="hpu" if is_hpex_available() else "cuda"
106
+ # in case that model is set to cuda:0 by default
107
+ if args.device_map.isdigit() and device=="cuda":
108
+ device = f"{device}:{args.device_map}"
109
110
if args.quantize:
111
if args.dtype in ["uNVFP4", "NVFP4+"]:
0 commit comments