Update app.py

app.py

@@ -2,13 +2,13 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import pickle
 
 
-
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 checkpoint = "my_t5.sav"
 
 #load model from drive
 with open(checkpoint, "rb") as f:
-    model=  pickle.load(f)
-
+    model=  pickle.load(f, map_location='cpu')
+    
 #tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 #model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)