diff --git a/tests/test_prediction.py b/tests/test_prediction.py
index 996ca1770bd291da5778e61c3c051ee982265180..6be77e9226ebc54a16c6888e7d1b6e779970b469 100644
--- a/tests/test_prediction.py
+++ b/tests/test_prediction.py
@@ -3,6 +3,7 @@
 import json
 import shutil
 
+import numpy as np
 import pytest
 import yaml