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