From 5b26a64d483e0025c42ab4a6955e14613771f452 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Thu, 8 Dec 2022 21:55:45 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 0fe89d5c..7de3c2bb 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -23,7 +23,7 @@ from torch import Tensor import torch.utils.checkpoint as checkpoint from .features import FeatureInfo -from .fx_features import register_notrace_function +from .fx_features import register_notrace_function, register_notrace_module from .helpers import build_model_with_cfg, pretrained_cfg_for_features from .layers import DropPath, to_2tuple, trunc_normal_, SelectAdaptivePool2d, ClassifierHead, Mlp from .pretrained import generate_default_cfgs