Test resnet classifier on CUDA

malfet · web-flow · commit b44971de83a6 · 2022-10-13T09:27:41.000-07:00
diff --git a/test/smoke_test/smoke_test.py b/test/smoke_test/smoke_test.py
@@ -195,6 +195,8 @@ def main() -> None:
     smoke_test_torchvision()
     smoke_test_torchvision_read_decode()
     smoke_test_torchvision_resnet50_classify()
+    if torch.cuda.is_available():
+        smoke_test_torchvision_resnet50_classify("cuda")
 
 
 if __name__ == "__main__":
