Option to run resnet classifier on specific device

pytorch · Oct 13, 2022 · 3b93537 · 3b93537
1 parent 6bc0bc2
commit 3b93537
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/test/smoke_test/smoke_test.py b/test/smoke_test/smoke_test.py
@@ -138,15 +138,15 @@ def smoke_test_torchvision_read_decode() -> None:
         raise RuntimeError(f"Unexpected shape of img_png: {img_png.shape}")
 
 
-def smoke_test_torchvision_resnet50_classify() -> None:
+def smoke_test_torchvision_resnet50_classify(device: str = "cpu") -> None:
     from torchvision.io import read_image
     from torchvision.models import resnet50, ResNet50_Weights
 
-    img = read_image(str(SCRIPT_DIR / "assets" / "dog2.jpg"))
+    img = read_image(str(SCRIPT_DIR / "assets" / "dog2.jpg")).to(device)
 
     # Step 1: Initialize model with the best available weights
     weights = ResNet50_Weights.DEFAULT
-    model = resnet50(weights=weights)
+    model = resnet50(weights=weights).to(device)
     model.eval()
 
     # Step 2: Initialize the inference transforms