From 4a54ed4d26f13f02c8fd9488da72d3b2805e7254 Mon Sep 17 00:00:00 2001
From: wuyefeilin <30919197+wuyefeilin@users.noreply.github.com>
Date: Mon, 9 Dec 2019 20:47:44 +0800
Subject: [PATCH] update solver.py and model_builder.py (#110)

* update solver.py and model_builder.py

* update solver.py
---
 pdseg/models/model_builder.py | 2 +-
 pdseg/solver.py               | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/pdseg/models/model_builder.py b/pdseg/models/model_builder.py
index 56a959a9..fd74da4e 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 0e3b1c57..baec69be 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):
-- 
GitLab