feat: map location

app.py

@@ -42,7 +42,10 @@ def get_model():
         num_hidden_layers=12,
     )
     model.load_state_dict(
-        torch.load(os.path.join("checkpoints", "PLAIN_ema_0.9999_360000.pt"))
+        torch.load(
+            os.path.join("checkpoints", "PLAIN_ema_0.9999_360000.pt"),
+            map_location=torch.device("cpu"),
+        )
     )
     model.eval()
     return model