!10440 fix gpu deepfm mindir load failed

From: @zhupuxu Reviewed-by: @zhoufeng54,@kisnwang Signed-off-by: @kisnwang
4 years ago · 19ef45f1a4
parent 1ae2569739 33dc7d3755
commit 19ef45f1a4
1 changed files with 1 additions and 0 deletions
--- a/model_zoo/official/recommend/deepfm/export.py
+++ b/model_zoo/official/recommend/deepfm/export.py
@ -39,6 +39,7 @@ if __name__ == "__main__":

    model_builder = ModelBuilder(ModelConfig, TrainConfig)
    _, network = model_builder.get_train_eval_net()
+    network.set_train(False)

    load_checkpoint(args.ckpt_file, net=network)