Satyajithchary commited on
Commit
0e13941
1 Parent(s): 010415d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -12,7 +12,12 @@ from PIL import Image
12
  from segment_anything_2 import SAM2ImagePredictor, build_sam2
13
 
14
  # Load your model
15
- device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 
 
 
16
  checkpoint = "checkpoints/sam2_hiera_large.pt"
17
  model_cfg = "sam2_hiera_l.yaml"
18
  model = build_sam2(model_cfg, checkpoint, device=device)
 
12
  from segment_anything_2 import SAM2ImagePredictor, build_sam2
13
 
14
  # Load your model
15
+ import torch
16
+
17
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
18
+ model = model.to(device)
19
+
20
+ #device = "cuda" if torch.cuda.is_available() else "cpu"
21
  checkpoint = "checkpoints/sam2_hiera_large.pt"
22
  model_cfg = "sam2_hiera_l.yaml"
23
  model = build_sam2(model_cfg, checkpoint, device=device)