diff --git a/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py b/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py
index dab485063ab2..e52889e3be6c 100644
--- a/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py
+++ b/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py
@@ -10,6 +10,7 @@
 import random
 import numpy as np
 import inspect
+import pytest
 
 MANUAL_SEED = 0
 random.seed(MANUAL_SEED)