Update model.py

layumi · Jul 2, 2024 · 9a92dd4 · 9a92dd4
1 parent d6304b6
commit 9a92dd4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model.py b/model.py
@@ -57,7 +57,7 @@ def __init__(self, input_dim, class_num, droprate, relu=False, bnorm=True, linea
         classifier = nn.Sequential(*classifier)
         classifier.apply(weights_init_classifier)
 
-        classifier.linear_num = linear
+        self.linear_num = linear
         self.add_block = add_block
         self.classifier = classifier
     def forward(self, x):