omg model loading is going to kill me torch 'just do it'

app.py

@@ -25,10 +25,13 @@ 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',
-                      map_location=torch.device(DEVICE)))
+model = NewCRFDepth(version='large07', inv_depth=False,
+                    max_depth=10).to(DEVICE).eval()
+model = torch.nn.DataParallel(model)
+model.load_state_dict()
+checkpoint = torch.load('checkpoints/nyu_L.pth',
+                        map_location=torch.device(DEVICE))
+model.load_state_dict(checkpoint['model'])
 
 title = "# IEBins: Iterative Elastic Bins for Monocular Depth Estimation"
 description = """Demo for **IEBins: Iterative Elastic Bins for Monocular Depth Estimation**.