Update README.md

README.md

@@ -46,11 +46,13 @@ 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, there are eight total
+)  # shards are ~4GB each, there are eight total
 
 prompt = "I was wondering how much wood a woodchuck could chuck? <answer>"
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-outputs = model.generate(**inputs, max_new_tokens=300) # default generation config (+ 300 tokens)
+outputs = model.generate(
+    **inputs, max_new_tokens=300
+)  # default generation config (+ 300 tokens)
 result = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
 result = result.split("<end_answer>")[0].strip()