diff --git a/video_prediction_tools/env_setup/wrapper_container.sh b/video_prediction_tools/env_setup/wrapper_container.sh
index fea29a0a9018a5436122389164cfff0859f22552..cfe716bee9f610b4a44988fc2ff6e4be048d06b4 100755
--- a/video_prediction_tools/env_setup/wrapper_container.sh
+++ b/video_prediction_tools/env_setup/wrapper_container.sh
@@ -27,6 +27,8 @@ export PYTHONPATH=/usr/local/lib/python3.8/dist-packages:$PYTHONPATH
 # ... and modules from this project
 export PYTHONPATH=${WORKING_DIR}:$PYTHONPATH
 export PYTHONPATH=${WORKING_DIR}/utils:$PYTHONPATH
+export PYTHONPATH=${WORKING_DIR}/model_modules:$PYTHONPATH
+export PYTHONPATH=${WORKING_DIR}/postprocess:$PYTHONPATH
 
 # Control
 echo "****** Check PYTHONPATH *****"