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