From 9560b01fec359bbbe7e82a30f6512c2d9febebeb Mon Sep 17 00:00:00 2001 From: randaller Date: Sun, 19 Mar 2023 19:31:09 +0300 Subject: [PATCH] Update hf-inference-cuda-example.py --- hf-inference-cuda-example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hf-inference-cuda-example.py b/hf-inference-cuda-example.py index 857ca17..679b9bd 100644 --- a/hf-inference-cuda-example.py +++ b/hf-inference-cuda-example.py @@ -22,4 +22,4 @@ device_map = infer_auto_device_map(model, max_memory={0: "6GiB", "cpu": "128GiB" print(device_map) batch = tokenizer("The highest mountain in China is ", return_tensors="pt") -print(tokenizer.decode(model.generate(batch["input_ids"].cuda(), max_length=100)[0])) +print(tokenizer.decode(model.generate(batch["input_ids"].cpu(), do_sample=True, top_k=50, max_length=100, top_p=0.95, temperature=1.0)[0]))