Spaces:

umuthopeyildirim
/

IEBins-Depth-Estimation

Sleeping

umuthopeyildirim commited on Jan 22

Commit

7ab01ae

•

1 Parent(s): 7b4afc3

Update model loading to support both CPU and GPU

Files changed (1) hide show

app.py CHANGED Viewed

@@ -27,7 +27,8 @@ css = """
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = NewCRFDepth(version="large07", inv_depth=False,
                     max_depth=10, pretrained=None).to(DEVICE).eval()
-model.load_state_dict(torch.load('checkpoints/nyu_L.pth'))
 title = "# IEBins: Iterative Elastic Bins for Monocular Depth Estimation"
 description = """Demo for **IEBins: Iterative Elastic Bins for Monocular Depth Estimation**.

 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = NewCRFDepth(version="large07", inv_depth=False,
                     max_depth=10, pretrained=None).to(DEVICE).eval()
+model.load_state_dict(torch.load('checkpoints/nyu_L.pth',
+                      map_location=torch.device(DEVICE)))
 title = "# IEBins: Iterative Elastic Bins for Monocular Depth Estimation"
 description = """Demo for **IEBins: Iterative Elastic Bins for Monocular Depth Estimation**.