diff --git a/py/prediction.py b/py/prediction.py
index 2ce7bee..e099b59 100644
--- a/py/prediction.py
+++ b/py/prediction.py
@@ -12,6 +12,7 @@
 from monai.transforms import AsDiscrete
 from monai.data import decollate_batch
 import shutil
+import vtk
 from scipy import linalg
 
 def main(args):