Remove redundant num_features arg from Classification model (#88)

* remove num_features arg from Classification model * add annotations to args Co-authored-by: Jirka Borovec <[email protected]>
Lightning-Universe · Feb 8, 2021 · 841986a · 841986a
1 parent 13acf78
commit 841986a
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/flash/vision/classification/model.py b/flash/vision/classification/model.py
@@ -38,10 +38,9 @@ class ImageClassifier(ClassificationTask):
 
     def __init__(
         self,
-        num_classes,
-        backbone="resnet18",
-        num_features: int = None,
-        pretrained=True,
+        num_classes: int,
+        backbone: str = "resnet18",
+        pretrained: bool = True,
         loss_fn: Callable = F.cross_entropy,
         optimizer: Type[torch.optim.Optimizer] = torch.optim.SGD,
         metrics: Union[Callable, Mapping, Sequence, None] = (Accuracy()),