From ebfcc130bd4727fe8c8437d502aa316936773df5 Mon Sep 17 00:00:00 2001
From: Bin Lu <lubin10@baidu.com>
Date: Fri, 10 Sep 2021 11:45:09 +0800
Subject: [PATCH] Update engine.py

---
 ppcls/engine/engine.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ppcls/engine/engine.py b/ppcls/engine/engine.py
index 34edc175..e43b198f 100644
--- a/ppcls/engine/engine.py
+++ b/ppcls/engine/engine.py
@@ -378,6 +378,7 @@ class ExportModel(nn.Layer):
         self.infer_output_key = config.get("infer_output_key", None)
         if self.infer_output_key == "features" and isinstance(self.base_model,
                                                               RecModel):
+            self.base_model.head = IdentityHead()
         if config.get("infer_add_softmax", True):
             self.softmax = nn.Softmax(axis=-1)
         else:
-- 
GitLab