Spaces:

stamps-labs
/

stamp2vec

Runtime error

sadjava commited on Jul 9, 2023

Commit

05aeeac

•

1 Parent(s): 1f84b9c

Fix for cpu

Files changed (1) hide show

app.py CHANGED Viewed

@@ -18,7 +18,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 yolov8 = YOLO(hf_hub_download('stamps-labs/yolov8-finetuned', filename='best.torchscript'), task='detect')
 yolo_stamp = YOLOStamp()
-yolo_stamp.load_state_dict(torch.load(hf_hub_download('stamps-labs/yolo-stamp', filename='state_dict.pth')))
 yolo_stamp = yolo_stamp.to(device)
 yolo_stamp.eval()
 transform = A.Compose([
@@ -31,7 +31,7 @@ vits8 = vits8.to(device)
 vits8.eval()
 encoder = Encoder()
-encoder.load_state_dict(torch.load(hf_hub_download('stamps-labs/vae-encoder', filename='encoder.pth')))
 encoder = encoder.to(device)
 encoder.eval()

 yolov8 = YOLO(hf_hub_download('stamps-labs/yolov8-finetuned', filename='best.torchscript'), task='detect')
 yolo_stamp = YOLOStamp()
+yolo_stamp.load_state_dict(torch.load(hf_hub_download('stamps-labs/yolo-stamp', filename='state_dict.pth'), map_location='cpu'))
 yolo_stamp = yolo_stamp.to(device)
 yolo_stamp.eval()
 transform = A.Compose([
 vits8.eval()
 encoder = Encoder()
+encoder.load_state_dict(torch.load(hf_hub_download('stamps-labs/vae-encoder', filename='encoder.pth'), map_location='cpu'))
 encoder = encoder.to(device)
 encoder.eval()