diff --git a/pdseg/models/model_builder.py b/pdseg/models/model_builder.py
index 56a959a9c20e1322ee5e906008e8bdcf392bd044..fd74da4e7eaf484a6d5d1ef110d1e6f7ce4eab0a 100644
--- a/pdseg/models/model_builder.py
+++ b/pdseg/models/model_builder.py
@@ -146,7 +146,7 @@ def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
             if ModelPhase.is_predict(phase):
                 origin_image = fluid.layers.data(
                     name='image',
-                    shape=[-1, 1, 1, cfg.DATASET.DATA_DIM],
+                    shape=[-1, -1, -1, cfg.DATASET.DATA_DIM],
                     dtype='float32',
                     append_batch_size=False)
                 image = fluid.layers.transpose(origin_image, [0, 3, 1, 2])
diff --git a/pdseg/solver.py b/pdseg/solver.py
index 0e3b1c57f774b050c0804f3149ad7301cdfa5ff3..baec69be06c4c2eb6d5e78480d5ffe9e1c737434 100644
--- a/pdseg/solver.py
+++ b/pdseg/solver.py
@@ -18,7 +18,10 @@ import paddle.fluid as fluid
 import numpy as np
 import importlib
 from utils.config import cfg
-from paddle.fluid.contrib.mixed_precision.decorator import OptimizerWithMixedPrecison, decorate, AutoMixedPrecisionLists
+try:
+    from paddle.fluid.contrib.mixed_precision.decorator import OptimizerWithMixedPrecison, decorate, AutoMixedPrecisionLists
+except:
+    from paddle.fluid.contrib.mixed_precision.decorator import OptimizerWithMixedPrecision, decorate, AutoMixedPrecisionLists
 
 
 class Solver(object):