diff --git a/gpu_mnist_example/train_gcu.py b/gpu_mnist_example/train_gcu.py
index c9e97f6..2dba5f0 100644
--- a/gpu_mnist_example/train_gcu.py
+++ b/gpu_mnist_example/train_gcu.py
@@ -47,6 +47,8 @@ import argparse
 import os
 os.system("pip install {}".format(os.getenv("OPENI_SDK_PATH")))
 
+from openi.context import prepare, upload_openi
+
 import importlib.util
 
 def is_torch_dtu_available():
@@ -80,6 +82,7 @@ if __name__ == '__main__':
     print("dataset_path:")
     print(os.listdir(dataset_path))
     os.listdir(dataset_path)
+    
     print("pretrain_model_path:")
     print(os.listdir(pretrain_model_path))
     os.listdir(pretrain_model_path)