nvidia
/

Mistral-NeMo-Minitron-8B-Base

Text Generation

text-generation-inference

Model card Files Files and versions

srvm commited on Aug 22, 2024

Commit

6628063

·

1 Parent(s): f5c6c67

Fix usage example

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -61,7 +61,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
 device = 'cuda'
 dtype = torch.bfloat16
-model = LlamaForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
 # Prepare the input text
 prompt = 'Complete the paragraph: our solar system is'

 device = 'cuda'
 dtype = torch.bfloat16
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
 # Prepare the input text
 prompt = 'Complete the paragraph: our solar system is'