From ca2b51d138586e707e0419ae2b42c8249b7f98e1 Mon Sep 17 00:00:00 2001
From: lixingjian <1762778193@qq.com>
Date: Fri, 17 Jan 2020 13:42:39 +0800
Subject: [PATCH] Update autodl.py

add parameter class_dim
---
 PaddleCV/image_classification/models/autodl.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/PaddleCV/image_classification/models/autodl.py b/PaddleCV/image_classification/models/autodl.py
index 915bc163..0c23d5b1 100644
--- a/PaddleCV/image_classification/models/autodl.py
+++ b/PaddleCV/image_classification/models/autodl.py
@@ -505,9 +505,9 @@ def StemConv1(input, C_out):
     return bn_a
 
 class NetworkImageNet(object):
-    def __init__(self, arch='DARTS_6M'):
+    def __init__(self, arch='DARTS_6M', class_dim=1000):
         self.params = train_parameters
-        self.class_num = 1000
+        self.class_num = class_dim
         self.init_channel = 48
         self._layers = 14
         self._auxiliary = False
-- 
GitLab