Remove torch.jit calls except for .ignore and .is_scripting to avoid 2.11 deprecation. Fix #2663

Author: rwightman

tests/test_models.py

@@ -82,7 +82,7 @@
     EXCLUDE_FILTERS = ['*enormous*', '*_7b_*']
     NON_STD_EXCLUDE_FILTERS = ['*gigantic*', '*enormous*', '*_3b_*', '*_7b_*']
 
-EXCLUDE_JIT_FILTERS = ['hiera_*', '*naflex*', '*_7b_*']
+EXCLUDE_JIT_FILTERS = ['hiera_*', '*naflex*', '*_7b_*', 'hrnet*', 'dpn*', 'densenet*', 'selecsls*']
 
 TARGET_FWD_SIZE = MAX_FWD_SIZE = 384
 TARGET_BWD_SIZE = 128

timm/models/densenet.py

@@ -72,24 +72,13 @@ def any_requires_grad(self, x: List[torch.Tensor]) -> bool:
                 return True
         return False
 
-    @torch.jit.unused  # noqa: T484
     def call_checkpoint_bottleneck(self, x: List[torch.Tensor]) -> torch.Tensor:
         """Call bottleneck function with gradient checkpointing."""
         def closure(*xs):
             return self.bottleneck_fn(xs)
 
         return checkpoint(closure, *x)
 
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (List[torch.Tensor]) -> (torch.Tensor)
-        pass
-
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (torch.Tensor) -> (torch.Tensor)
-        pass
-
     # torchscript does not yet support *args, so we overload method
     # allowing it to take either a List[Tensor] or single Tensor
     def forward(self, x: Union[torch.Tensor, List[torch.Tensor]]) -> torch.Tensor:  # noqa: F811

timm/models/dpn.py

@@ -34,16 +34,6 @@ def __init__(
         super().__init__()
         self.bn = norm_layer(in_chs, eps=0.001, **dd)
 
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (Tuple[torch.Tensor, torch.Tensor]) -> (torch.Tensor)
-        pass
-
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (torch.Tensor) -> (torch.Tensor)
-        pass
-
     def forward(self, x):
         if isinstance(x, tuple):
             x = torch.cat(x, dim=1)
@@ -124,16 +114,6 @@ def __init__(
             self.c1x1_c1 = None
             self.c1x1_c2 = None
 
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]
-        pass
-
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
-        pass
-
     def forward(self, x) -> Tuple[torch.Tensor, torch.Tensor]:
         if isinstance(x, tuple):
             x_in = torch.cat(x, dim=1)

timm/models/hrnet.py

@@ -502,28 +502,12 @@ class SequentialList(nn.Sequential):
     def __init__(self, *args):
         super().__init__(*args)
 
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (List[torch.Tensor]) -> (List[torch.Tensor])
-        pass
-
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (torch.Tensor) -> (List[torch.Tensor])
-        pass
-
     def forward(self, x) -> List[torch.Tensor]:
         for module in self:
             x = module(x)
         return x
 
 
-@torch.jit.interface
-class ModuleInterface(torch.nn.Module):
-    def forward(self, input: torch.Tensor) -> torch.Tensor: # `input` has a same name in Sequential forward
-        pass
-
-
 block_types_dict = {
     'BASIC': BasicBlock,
     'BOTTLENECK': Bottleneck
@@ -816,7 +800,7 @@ def forward_features(self, x):
             if y is None:
                 y = incre(yl[i])
             else:
-                down: ModuleInterface = self.downsamp_modules[i - 1]  # needed for torchscript module indexing
+                down = self.downsamp_modules[i - 1]
                 y = incre(yl[i]) + down.forward(y)
 
         y = self.final_layer(y)

timm/models/selecsls.py

@@ -27,16 +27,6 @@ class SequentialList(nn.Sequential):
     def __init__(self, *args):
         super().__init__(*args)
 
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (List[torch.Tensor]) -> (List[torch.Tensor])
-        pass
-
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (torch.Tensor) -> (List[torch.Tensor])
-        pass
-
     def forward(self, x) -> List[torch.Tensor]:
         for module in self:
             x = module(x)
@@ -49,16 +39,6 @@ def __init__(self, mode='index', index=0):
         self.mode = mode
         self.index = index
 
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (List[torch.Tensor]) -> (torch.Tensor)
-        pass
-
-    @torch.jit._overload_method  # noqa: F811
-    def forward(self, x):
-        # type: (Tuple[torch.Tensor]) -> (torch.Tensor)
-        pass
-
     def forward(self, x) -> torch.Tensor:
         if self.mode == 'index':
             return x[self.index]