diff --git a/video_prediction_tools/model_modules/video_prediction/models/base_model.py b/video_prediction_tools/model_modules/video_prediction/models/base_model.py
index 2bc8a399a49a10f3df99f9646c040140970d573c..1857f8b915d62646dff9a73d63f16e78656ddc57 100644
--- a/video_prediction_tools/model_modules/video_prediction/models/base_model.py
+++ b/video_prediction_tools/model_modules/video_prediction/models/base_model.py
@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: MIT
 
+import functools
 import itertools
 import os
 import re