Update engine.py

2021-09-10 11:45:09 +08:00 · 2021-09-10 11:45:09 +08:00 · ebfcc130bd
parent 2ab10a2d95
commit ebfcc130bd
1 changed files with 1 additions and 0 deletions
--- 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: