Update app.py

app.py

@@ -2,11 +2,24 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import pickle
 
 
+import io
+
+class CPU_Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == 'torch.storage' and name == '_load_from_bytes':
+            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
+        else:
+        return super().find_class(module, name)
+
+#contents = pickle.load(f) becomes...
+#contents = CPU_Unpickler(f).load()
+
+
 checkpoint = "my_t5.sav"
 
 #load model from drive
 with open(checkpoint, "rb") as f:
-    model=  pickle.load(f)
+    model=  CPU_Unpickler(f).load()
     
 #tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 #model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)