未验证 提交 ebfcc130 编写于 作者: B Bin Lu 提交者: GitHub

Update engine.py

上级 2ab10a2d
...@@ -378,6 +378,7 @@ class ExportModel(nn.Layer): ...@@ -378,6 +378,7 @@ class ExportModel(nn.Layer):
self.infer_output_key = config.get("infer_output_key", None) self.infer_output_key = config.get("infer_output_key", None)
if self.infer_output_key == "features" and isinstance(self.base_model, if self.infer_output_key == "features" and isinstance(self.base_model,
RecModel): RecModel):
self.base_model.head = IdentityHead()
if config.get("infer_add_softmax", True): if config.get("infer_add_softmax", True):
self.softmax = nn.Softmax(axis=-1) self.softmax = nn.Softmax(axis=-1)
else: else:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册