Update README.md

README.md

@@ -6,6 +6,7 @@ config version: 1
 torchao version: 0.14.dev
 ```
 
+# Generate Quantized Model
 ```
 import logging
 
@@ -106,4 +107,54 @@ output_text = tokenizer.batch_decode(
 print("Response:", output_text[0][len(prompt) :])
 
 assert(correct_output_text == output_text)
+```
+
+
+# Test Loading
+```
+from transformers import (
+  AutoModelForCausalLM,
+  AutoProcessor,
+  AutoTokenizer,
+  TorchAoConfig,
+)
+from torchao.quantization import Float8Tensor
+from torchao.quantization import (
+    Float8Tensor,
+    Int4TilePackedTo4dTensor,
+    IntxUnpackedToInt8Tensor,
+)
+import torch
+
+model_name = "torchao-testing/opt-125m-ModuleFqnToConfig-v1-regex-0.14.0.dev"
+device = "cuda"
+input_text = "What are we having for dinner?"
+max_new_tokens = 10
+
+quantized_model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map=device,
+    dtype=torch.bfloat16,
+)
+print("quantized model:", quantized_model)
+for i in range(12):
+    if i == 3:
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Int4TilePackedTo4dTensor)
+    else:
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Float8Tensor)
+    assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Float8Tensor)
+    assert not isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Float8Tensor)
+    assert isinstance(quantized_model.model.decoder.layers[i].self_attn.out_proj.weight, IntxUnpackedToInt8Tensor)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+input_ids = tokenizer(input_text, return_tensors="pt").to(device)
+
+output = quantized_model.generate(**input_ids, max_new_tokens=max_new_tokens)
+EXPECTED_OUTPUT = [
+    "What are we having for dinner?\n\nJessica: (smiling)",
+    "What are we having for dinner?\n\nJess: (smiling) I",
+]
+# self.assertTrue(tokenizer.decode(output[0], skip_special_tokens=True) in EXPECTED_OUTPUT)
+
 ```