Fixed fp32 benchmark test failure for PyTorch ipex backend
@@ -313,6 +313,14 @@ def main_worker(gpu, ngpus_per_node, args):
from lpot.utils.pytorch import load
new_model = load(
os.path.abspath(os.path.expanduser(args.tuned_checkpoint)), model)
+ else:
+ if args.ipex:
+ # TODO: It will remove when IPEX spport to save script model.
+ model.to(ipex.DEVICE)
+ try:
+ new_model = torch.jit.script(model)
+ except:
+ new_model = torch.jit.trace(model, torch.randn(1, 3, 224, 224).to(ipex.DEVICE))
else:
model.fuse_model()
new_model = model
