Update app.py

app.py

@@ -5,7 +5,12 @@ import torch
 from segment_anything import sam_model_registry
 from automatic_mask_generator import SamAutomaticMaskGenerator
 
-device = 'cuda'
+
+if torch.cuda.is_available():
+    device = 'cuda'
+else:
+    device = 'cpu'
+
 sam = sam_model_registry['vit_h'](checkpoint='./sam_vit_h_4b8939.pth')
 sam.to(device=device)