Fixed default num_classes to 10 for CIFAR10 (#1176)

pytorch · Jul 28, 2023 · 741de70 · 741de70
1 parent 6fc19c7
commit 741de70
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vision_transformer/main.py b/vision_transformer/main.py
@@ -225,7 +225,7 @@ def main():
                         help='dropout value (default : 0.1)')
     parser.add_argument('--img-size', type=int, default=224,
                         help='image size to be reshaped to (default : 224')
-    parser.add_argument('--num-classes', type=int, default=16,
+    parser.add_argument('--num-classes', type=int, default=10,
                         help='number of classes in dataset (default : 10 for CIFAR10)')
     parser.add_argument('--epochs', type=int, default=10,
                         help='number of epochs (default : 10)')