From e8ba9118d9475fa47b4d7108c16b7e166087b193 Mon Sep 17 00:00:00 2001
From: Nikita Shulga <nshulga@fb.com>
Date: Thu, 13 Oct 2022 09:27:41 -0700
Subject: [PATCH] Test resnet classifier on CUDA

---
 test/smoke_test/smoke_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/test/smoke_test/smoke_test.py b/test/smoke_test/smoke_test.py
index 157170975..e069a712e 100644
--- 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__":