diff --git a/tests/test_prediction.py b/tests/test_prediction.py
index bb201d3a4f642b7144f865eb3b32e8e229f2f342..fbe69846ee0c3808f8efb3b9ae9bc891d25ffd86 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