Update README.md

README.md

@@ -31,7 +31,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 tokenizer = AutoTokenizer.from_pretrained("rudyon/rudygpt")
 model = GPT(GPTConfig(depth=12, vocab_size=50304))
-state_dict = torch.load("pytorch_model.bin", map_location='cpu')
+state_dict = torch.load("pytorch_model.bin", map_location=device)
 model.load_state_dict(state_dict)
 model.eval()
 print(model.generate("Hello!"))