Update app.py

app.py

@@ -17,9 +17,18 @@ tokenizer = AutoTokenizer.from_pretrained('t5-small')
 #model = torch.load(model_name+'/model.pt')
 
 #if model_name == "T5 Small":
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, checkpoint_file="t5_epoch9.ckpt")
+#model = AutoModelForSeq2SeqLM.from_pretrained(model_name, checkpoint_file="t5_epoch9.ckpt")
     #model = torch.load(model_name+'/model.pt')
 checkpoint_path=model_name+'/t5_epoch9.ckpt'
+
+
+# Choose the appropriate device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model_state_dict = torch.load(checkpoint_path, map_location=device)
+model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")
+model.load_state_dict(model_state_dict)
+
+
 #model=T5.load_from_checkpoint(checkpoint_path, map_location=torch.device('cpu'))
     
 #else: