Addressing comments

Author: atalman

test/smoke_test/smoke_test.py

@@ -3,11 +3,12 @@
 import torch
 import torchvision
 import torchaudio
-import pathlib
+from pathlib import Path
 
 gpu_arch_ver = os.getenv('GPU_ARCH_VER')
 gpu_arch_type = os.getenv('GPU_ARCH_TYPE')
 is_cuda_system = gpu_arch_type == 'cuda'
+SCRIPT_DIR = Path(__file__).parent
 
 def smoke_test_cuda() -> None:
     if(not torch.cuda.is_available() and is_cuda_system):
@@ -42,16 +43,18 @@ def smoke_test_torchvision() -> None:
 
 def smoke_test_torchvision_read_decode() -> None:
     from torchvision.io import read_image
-    img_jpg = read_image(str(pathlib.Path(__file__).parent / 'assets' / 'rgb_pytorch.jpg'))
-    assert img_jpg.ndim == 3 and img_jpg.numel() > 100
-    img_png = read_image(str(pathlib.Path(__file__).parent  / 'assets' / 'rgb_pytorch.png'))
-    assert img_png.ndim == 3 and img_png.numel() > 100
+    img_jpg = read_image(str(SCRIPT_DIR / 'assets' / 'rgb_pytorch.jpg'))
+    if img_jpg.ndim != 3 or img_jpg.numel() < 100:
+      raise RuntimeError(f'Unexpected shape of img_jpg: {img_jpg.shape}')
+    img_png = read_image(str(SCRIPT_DIR  / 'assets' / 'rgb_pytorch.png'))
+    if img_png.ndim != 3 or img_png.numel() < 100:
+      raise RuntimeError(f'Unexpected shape of img_png: {img_png.shape}')
 
 def smoke_test_torchvision_resnet50_classify() -> None:
     from torchvision.io import read_image
     from torchvision.models import resnet50, ResNet50_Weights
 
-    img = read_image(str(pathlib.Path(__file__).parent / 'assets' / 'dog2.jpg'))
+    img = read_image(str(SCRIPT_DIR / 'assets' / 'dog2.jpg'))
 
     # Step 1: Initialize model with the best available weights
     weights = ResNet50_Weights.DEFAULT