Update README.md

README.md

@@ -43,12 +43,12 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load Code Generator LLM and tokenizer from checkpoint
-tokenizer = AutoTokenizer.from_pretrained("DeathReaper0965/gpt2_large_code_generator/", )
-model = AutoModelForCausalLM.from_pretrained("DeathReaper0965/gpt2_large_code_generator/")
-
+tokenizer = AutoTokenizer.from_pretrained("DeathReaper0965/gpt2_large_code_generator")
+model = AutoModelForCausalLM.from_pretrained("DeathReaper0965/gpt2_large_code_generator")
 model = model.to("cuda" if torch.cuda.is_available() else "cpu")
 
-inputs = tokenizer("def hello_world():", return_tensors="pt").to("cuda")
+inputs = tokenizer("def hello_world():", return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+
 outputs = model.generate(**inputs,
                          max_new_tokens= 30,
                          temperature= 0.5,