hugohabicht01 commited on
Commit
973cd99
·
1 Parent(s): d715937
Files changed (1) hide show
  1. blurnonymize.py +2 -2
blurnonymize.py CHANGED
@@ -158,9 +158,9 @@ class ImageBlurnonymizer:
158
 
159
  # self.device = "cuda" if torch.cuda.is_available() else "cpu"
160
  # self.device = "cuda"
161
- device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
162
  # Set the device for PyTorch
163
- torch.cuda.set_device(device)
164
  self.predictor = SAM2ImagePredictor.from_pretrained(
165
  "facebook/sam2.1-hiera-small",
166
  device="cuda",
 
158
 
159
  # self.device = "cuda" if torch.cuda.is_available() else "cpu"
160
  # self.device = "cuda"
161
+ self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
162
  # Set the device for PyTorch
163
+ torch.cuda.set_device(self.device)
164
  self.predictor = SAM2ImagePredictor.from_pretrained(
165
  "facebook/sam2.1-hiera-small",
166
  device="cuda",