Skip to content

Commit

Permalink
fix code format (#7377)
Browse files Browse the repository at this point in the history
  • Loading branch information
pmeier authored Mar 2, 2023
1 parent caf12f8 commit 0e62c34
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion test/smoke_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,15 @@ def smoke_test_torchvision_read_decode() -> None:
if img_png.ndim != 3 or img_png.numel() < 100:
raise RuntimeError(f"Unexpected shape of img_png: {img_png.shape}")


def smoke_test_compile() -> None:
model = resnet50().cuda()
model = torch.compile(model)
x = torch.randn(1, 3, 224, 224, device="cuda")
out = model(x)
print(f"torch.compile model output: {out.shape}")


def smoke_test_torchvision_resnet50_classify(device: str = "cpu") -> None:
img = read_image(str(SCRIPT_DIR / ".." / "gallery" / "assets" / "dog2.jpg")).to(device)

Expand Down Expand Up @@ -73,6 +75,5 @@ def main() -> None:
smoke_test_torchvision_resnet50_classify("mps")



if __name__ == "__main__":
main()

0 comments on commit 0e62c34

Please sign in to comment.