fix: load model

app.py

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