diff --git a/tests/test_visualization.py b/tests/test_visualization.py
index 4eb51ee..e66f27b 100644
--- a/tests/test_visualization.py
+++ b/tests/test_visualization.py
@@ -6,7 +6,7 @@ import torch
 
 # Import visualization functions
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
-from scripts.visualize_predictions import visualize_prediction
+from scripts.visualize_predictions import visualize_prediction  # noqa: E402
 
 
 def test_visualize_prediction():