Spaces:
Running
Running
testing autocast to fp16
Browse files
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 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
|
|
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 |
|