joselobenitezg commited on
Commit
a92daf2
·
1 Parent(s): 9eb921d

update seg inference

Browse files
Files changed (1) hide show
  1. inference/seg.py +3 -0
inference/seg.py CHANGED
@@ -16,6 +16,9 @@ def load_model(task, version):
16
  return None, None
17
 
18
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 
19
  model = torch.jit.load(model_path)
20
  model.eval()
21
  model.to(device)
 
16
  return None, None
17
 
18
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
19
+ if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
20
+ torch.backends.cuda.matmul.allow_tf32 = True
21
+ torch.backends.cudnn.allow_tf32 = True
22
  model = torch.jit.load(model_path)
23
  model.eval()
24
  model.to(device)