Update app.py

app.py

@@ -195,7 +195,7 @@ with tempfile.TemporaryDirectory() as tmpdir:
 
             dam2 = get_dam2_model()
             depth_dam2 = dam2.infer_image(image_bgr)
-            depth = torch.tensor(depth_dam2).unsqueeze(0).unsqueeze(0).float().cuda()
+            depth = torch.tensor(depth_dam2).unsqueeze(0).unsqueeze(0).float()
 
             depth_image = cv2.applyColorMap((normalize_disp(depth_dam2) * 255).astype(np.uint8), cv2.COLORMAP_JET)
 
@@ -203,7 +203,7 @@ with tempfile.TemporaryDirectory() as tmpdir:
 
         @spaces.GPU()
         def cb_generate(image, depth: Tensor, scale_factor):
-            norm_disp = normalize_disp(depth)
+            norm_disp = normalize_disp(depth.cuda())
             disp = norm_disp * scale_factor / 100 * IMAGE_SIZE
 
             genstereo = get_genstereo_model()