fix export pretrained_model when it is None

pull/1820/merge
HydrogenSulfate 2022-12-16 17:20:26 +08:00
parent 20725dd770
commit 97398d9f12
1 changed files with 9 additions and 6 deletions

View File

@ -548,12 +548,15 @@ class Engine(object):
"use_multilabel",
False) or "ATTRMetric" in self.config["Metric"]["Eval"][0]
model = ExportModel(self.config["Arch"], self.model, use_multilabel)
if self.config["Global"]["pretrained_model"].startswith("http"):
load_dygraph_pretrain_from_url(
model.base_model, self.config["Global"]["pretrained_model"])
else:
load_dygraph_pretrain(model.base_model,
self.config["Global"]["pretrained_model"])
if self.config["Global"]["pretrained_model"] is not None:
if self.config["Global"]["pretrained_model"].startswith("http"):
load_dygraph_pretrain_from_url(
model.base_model,
self.config["Global"]["pretrained_model"])
else:
load_dygraph_pretrain(
model.base_model,
self.config["Global"]["pretrained_model"])
model.eval()