Fix weight loading

app.py

@@ -54,7 +54,7 @@ class Ensemble(nn.Module):
 
 
 checkpoints = ["fold0.ckpt", "fold1.ckpt", "fold2.ckpt"]
-weights = [torch.load(ckpt)["state_dict"] for ckpt in checkpoints]
+weights = [torch.load(ckpt, map_location=torch.device("cpu"))["state_dict"] for ckpt in checkpoints]
 weights = [{k.replace("model.", "") : v for k, v in wt.items()} for wt in weights]
 models = [Net2D(wt) for wt in weights] 
 ensemble = Ensemble(models).eval()