mirrors
/
PaddleClas
mirror of https://github.com/PaddlePaddle/PaddleClas.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update engine.py

pull/1218/head
Bin Lu 2021-09-10 11:45:09 +08:00 committed by GitHub
parent 2ab10a2d95
commit ebfcc130bd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
ppcls/engine/engine.py
View File

@ -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: