Update README.md

README.md

@@ -39,12 +39,11 @@ tokenizer = AutoTokenizer.from_pretrained("pszemraj/pythia-6.9b-HC3")
 
 model = AutoModelForCausalLM.from_pretrained(
     "pszemraj/pythia-6.9b-HC3", load_in_8bit=True, device_map="auto"
-)
+) # shards are ~4GB each
 
 prompt = "I was wondering how much wood a woodchuck could chuck? <answer>"
-inputs = tokenizer(prompt, return_tensors="pt")
-inputs = inputs.to("cuda")
-outputs = model.generate(**inputs, max_new_tokens=300)
+inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+outputs = model.generate(**inputs, max_new_tokens=300) # default generation config (+ 300 tokens)
 result = tokenizer.batch_decode(outputs, skip_special_tokens=True)
 
 import pprint as pp