Spaces:

Manthanx
/

ImageSeg_IPPR

Runtime error

Manthanx commited on Jun 6, 2023

Commit

c836824

1 Parent(s): 3fab3aa

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -14,7 +14,7 @@ device = torch.device(device)
 model_path = 'cityscapes_dataUNet.pth'
 num_classes = 10
 model = UNet(num_classes=num_classes)
-model.load_state_dict(torch.load(model_path))
 model.to(device)
 model.eval()

 model_path = 'cityscapes_dataUNet.pth'
 num_classes = 10
 model = UNet(num_classes=num_classes)
+model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))
 model.to(device)
 model.eval()