Update app.py

app.py

@@ -43,9 +43,15 @@ def lr_schedule(current_step):
     return factor
 
 
+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)
+
 
 with open('learner.bin', 'rb') as f:
-    learner = pickle.load(f)
+    learner = CPU_Unpickler(f).load()
 
 
 def dg_predict(tweet):