jhj0517 commited on
Commit
7dc0c85
1 Parent(s): f3f4569

Update dtype for CPU

Browse files
Files changed (1) hide show
  1. modules/sam_inference.py +1 -1
modules/sam_inference.py CHANGED
@@ -41,7 +41,7 @@ class SamInference:
41
  self.model_dir = model_dir
42
  self.output_dir = output_dir
43
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
44
- self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
45
  self.mask_generator = None
46
  self.image_predictor = None
47
  self.video_predictor = None
 
41
  self.model_dir = model_dir
42
  self.output_dir = output_dir
43
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
44
+ self.dtype = torch.float16 if torch.cuda.is_available() else torch.bfloat16
45
  self.mask_generator = None
46
  self.image_predictor = None
47
  self.video_predictor = None