diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
index 01a215a293d3..c56fea05d604 100644
--- a/docker/Dockerfile.rocm
+++ b/docker/Dockerfile.rocm
@@ -400,6 +400,7 @@ def pytest_sessionfinish(session, exitstatus):
     _exit_code = int(exitstatus)
 
 def pytest_unconfigure(config):
+    import sys
     sys.stdout.flush()
     sys.stderr.flush()
     os._exit(_exit_code)