add annotations to args

Lightning-Universe · Feb 8, 2021 · 5f2936b · 5f2936b
1 parent 5acc96d
commit 5f2936b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/flash/vision/classification/model.py b/flash/vision/classification/model.py
@@ -38,9 +38,9 @@ class ImageClassifier(ClassificationTask):
 
     def __init__(
         self,
-        num_classes,
-        backbone="resnet18",
-        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()),