nvidia
/

Nemotron-Flash-3B

Text Generation

Model card Files Files and versions

YongganFu commited on 12 days ago

Commit

e623d13

·

verified ·

1 Parent(s): 972cbc6

Update README.md

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -93,7 +93,7 @@ from transformers import AutoConfig, AutoModelForCausalLM
 repo_name = "nvidia/Nemotron-Flash-3B"
 config = AutoConfig.from_pretrained(repo_name, trust_remote_code=True)
-setattr(config, "attention_implementation_new", "flash_attention_2")
 model = AutoModelForCausalLM.from_pretrained(repo_name, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True)
 ```

 repo_name = "nvidia/Nemotron-Flash-3B"
 config = AutoConfig.from_pretrained(repo_name, trust_remote_code=True)
+setattr(config, "attn_implementation_new", "flash_attention_2")
 model = AutoModelForCausalLM.from_pretrained(repo_name, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True)
 ```