nvidia
/

Nemotron-Flash-1B

Text Generation

Model card Files Files and versions

YongganFu commited on Jan 9

Commit

da87f70

·

verified ·

1 Parent(s): fb5d4e1

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-1B"
 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-1B"
 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)
 ```