David Vaillant commited on
Commit
7fe2d10
1 Parent(s): d804cc4
Files changed (1) hide show
  1. backend.py +1 -1
backend.py CHANGED
@@ -17,7 +17,7 @@ preproc = v2.Compose([
17
 
18
  # Load the necessary models.
19
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
20
- CHECKPOINT_FILE = os.getenv("SAM_FINETUNE_CHECKPOINT", "checkpoints/bbox_finetune.pth")
21
 
22
  processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
23
  tuned_model = SamModel.from_pretrained("facebook/sam-vit-large").to(device)
 
17
 
18
  # Load the necessary models.
19
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
20
+ CHECKPOINT_FILE = os.getenv("SAM_FINETUNE_CHECKPOINT", "checkpoints/bbox_finetune.ckpt")
21
 
22
  processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
23
  tuned_model = SamModel.from_pretrained("facebook/sam-vit-large").to(device)