set_train_false

Signed-off-by: zhupuxu <zhupuxu@huawei.com>
4 years ago · 33dc7d3755
parent 4df0baa6c7
commit 33dc7d3755
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)