!17280 fix deeptext export mindir error

From: @zhouneng2
Reviewed-by: @oacjiewen,@c_34
Signed-off-by: @c_34
This commit is contained in:
mindspore-ci-bot 2021-05-31 20:24:19 +08:00 committed by Gitee
commit 1f305914d0
1 changed files with 5 additions and 0 deletions

View File

@ -16,6 +16,7 @@
import numpy as np
import mindspore as ms
import mindspore.common.dtype as mstype
from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
from src.Deeptext.deeptext_vgg16 import Deeptext_VGG16_Infer
@ -40,6 +41,10 @@ if __name__ == '__main__':
load_param_into_net(net, param_dict_new)
device_type = "Ascend" if context.get_context("device_target") == "Ascend" else "Others"
if device_type == "Ascend":
net.to_float(mstype.float16)
img_data = Tensor(np.zeros([config.test_batch_size, 3, config.img_height, config.img_width]), ms.float32)
export(net, img_data, file_name=config.file_name, file_format=config.file_format)