Update README.md

README.md

@@ -28,6 +28,8 @@ prompt.
 
 ### Usage
 
+Inference on GPU with 4-bit quantization: 
+
 ```
 %pip install -qq transformers accelerate bitsandbytes
 ```
@@ -56,7 +58,7 @@ model = AutoModelForCausalLM.from_pretrained(
 ```python
 inputs = tokenizer("Once upon a time,", return_tensors='pt')
 inputs = {
-    k:v.cpu() for k,v in inputs.items()
+    k:v.cuda() for k,v in inputs.items()
 }
 outputs = model.generate(
     **inputs,
@@ -68,8 +70,8 @@ tokenizer.decode(outputs[0])
 ```
 
 TODO
-- test to see if model works with .from_pretrained <br>
-  - test fp32, fp16, 8 and 4 bit
+- ~~test to see if model works with .from_pretrained~~ <br>
+  - ~~test fp32, fp16, 8 and 4 bit~~
 - shard model to max 1gb for use in even lower vram settings <br>
 - safetensors <br>
 - upload bf16 version of model <br>