Update README.md

README.md

@@ -65,6 +65,10 @@ Extract address components from: {address}<|eot_id|><|start_header_id|>assistant
     # Tokenize
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
     
+    # FIX: Move inputs to the same device as the model
+    device = next(model.parameters()).device
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    
     # Generate
     with torch.no_grad():
         outputs = model.generate(