From 7aad9a73a25aa394f1b042e3740c996cb296419b Mon Sep 17 00:00:00 2001 From: Bai Yifan <me@ethanbai.com> Date: Fri, 18 Dec 2020 22:24:04 +0800 Subject: [PATCH] add dygraph demo issue (#562) --- demo/dygraph/quant/train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/demo/dygraph/quant/train.py b/demo/dygraph/quant/train.py index 25c0a3d1..a9853653 100644 --- a/demo/dygraph/quant/train.py +++ b/demo/dygraph/quant/train.py @@ -117,9 +117,9 @@ def compress(args): pretrain = True if args.data == "imagenet" else False if args.model == "mobilenet_v1": - net = mobilenet_v1(pretrained=pretrain) + net = mobilenet_v1(pretrained=pretrain, num_classes=class_dim) elif args.model == "mobilenet_v3": - net = MobileNetV3_large_x1_0() + net = MobileNetV3_large_x1_0(class_dim=class_dim) if pretrain: load_dygraph_pretrain(net, args.pretrained_model, True) else: -- GitLab