Update model.py

layumi · Jul 2, 2024 · db96c81 · db96c81
1 parent 9a66c11
commit db96c81
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/model.py b/model.py
@@ -57,6 +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.add_block = add_block
         self.classifier = classifier
     def forward(self, x):