vishesh-t27 commited on
Commit
31c8ace
·
verified ·
1 Parent(s): c5d242d

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +5 -3
README.md CHANGED
@@ -133,17 +133,19 @@ We finetuned our model and other open source models on [Google's IndicGenBench](
133
  !pip install transformers=='5.4.0'
134
 
135
  from transformers import AutoModelForCausalLM, AutoTokenizer
136
-
137
 
138
  model_name = "FrontiersMind/Nandi-mini-150M"
139
 
 
 
140
  tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
141
 
142
  model = AutoModelForCausalLM.from_pretrained(
143
  model_name,
144
  trust_remote_code=True,
145
- device_map="auto",
146
- ).eval()
147
 
148
 
149
  prompt = """
 
133
  !pip install transformers=='5.4.0'
134
 
135
  from transformers import AutoModelForCausalLM, AutoTokenizer
136
+ import torch
137
 
138
  model_name = "FrontiersMind/Nandi-mini-150M"
139
 
140
+ device = "cuda" if torch.cuda.is_available() else "cpu"
141
+
142
  tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
143
 
144
  model = AutoModelForCausalLM.from_pretrained(
145
  model_name,
146
  trust_remote_code=True,
147
+ dtype=torch.bfloat16
148
+ ).to(device).eval()
149
 
150
 
151
  prompt = """