ahsanMah commited on
Commit
0326791
·
verified ·
1 Parent(s): ff98280

testing autocast to fp16

Browse files
Files changed (1) hide show
  1. hfapp.py +5 -4
hfapp.py CHANGED
@@ -19,10 +19,11 @@ def run_inference(model, img):
19
  model = model.to('cuda')
20
  print("model on cuda:", next(model.scorenet.net.parameters()).is_cuda)
21
  print("img on cuda:", img.is_cuda)
22
- img = torch.nn.functional.interpolate(img, size=64, mode="bilinear")
23
- score_norms = model.scorenet(img)
24
- score_norms = score_norms.square().sum(dim=(2, 3, 4)) ** 0.5
25
- img_likelihood = model(img).cpu().numpy()
 
26
  score_norms = score_norms.cpu().numpy()
27
  return img_likelihood, score_norms
28
 
 
19
  model = model.to('cuda')
20
  print("model on cuda:", next(model.scorenet.net.parameters()).is_cuda)
21
  print("img on cuda:", img.is_cuda)
22
+ with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=True):
23
+ img = torch.nn.functional.interpolate(img, size=64, mode="bilinear")
24
+ score_norms = model.scorenet(img)
25
+ score_norms = score_norms.square().sum(dim=(2, 3, 4)) ** 0.5
26
+ img_likelihood = model(img).cpu().numpy()
27
  score_norms = score_norms.cpu().numpy()
28
  return img_likelihood, score_norms
29