DTPDC

Sleeping

sh0kul commited on Jun 8, 2023

Commit

c3826b2

1 Parent(s): 284a4a8

Fixing non CUDA error

Files changed (1) hide show

app.py CHANGED Viewed

@@ -99,7 +99,7 @@ class DefectResNet(pl.LightningModule):
 loadmodel = DefectResNet()
 def load_checkpoint(checkpoint):
     loadmodel.load_state_dict(checkpoint["state_dict"])
-load_checkpoint(torch.load("models/model.ckpt"))
 loadmodel.eval()
 transform = T.Compose([

 loadmodel = DefectResNet()
 def load_checkpoint(checkpoint):
     loadmodel.load_state_dict(checkpoint["state_dict"])
+load_checkpoint(torch.load("models/model.ckpt", map_location=torch.device('cpu')))
 loadmodel.eval()
 transform = T.Compose([