Update davit.py

3 years ago · d1e68fa33f
parent 4ba893ceae
commit d1e68fa33f
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -478,7 +478,7 @@ class DaViT(nn.Module):
        return self._features_only
        
    @features_only.setter
-    def features_only(self, new_value):
+    def features_only(self, new_value : bool):
        self._features_only = new_value
        self.forward = self._get_forward_fn()