diff --git a/tests/test_prediction.py b/tests/test_prediction.py
index 83317477247893b2886b3238cc908d79a00dda2d..37f860b6901266dba3988620333a1ea9fcda24ae 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