From 56c3a84db397fba690d704f0d07308c4214df2ce Mon Sep 17 00:00:00 2001 From: Jasha10 <8935917+Jasha10@users.noreply.github.com> Date: Fri, 22 Jul 2022 16:59:55 -0500 Subject: [PATCH] Update type hint for `register_notrace_module` register_notrace_module is used to decorate types (i.e. subclasses of nn.Module). It is not called on module instances. --- timm/models/fx_features.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/timm/models/fx_features.py b/timm/models/fx_features.py index a9c05b0a..b09381b7 100644 --- a/timm/models/fx_features.py +++ b/timm/models/fx_features.py @@ -1,7 +1,7 @@ """ PyTorch FX Based Feature Extraction Helpers Using https://pytorch.org/vision/stable/feature_extraction.html """ -from typing import Callable, List, Dict, Union +from typing import Callable, List, Dict, Union, Type import torch from torch import nn @@ -35,7 +35,7 @@ except ImportError: pass -def register_notrace_module(module: nn.Module): +def register_notrace_module(module: Type[nn.Module]): """ Any module not under timm.models.layers should get this decorator if we don't want to trace through it. """