From 64912a8c45a9d4dd50e89aecf566d50f6123405e Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Fri, 25 Aug 2023 13:33:13 -0700 Subject: [PATCH] Fix torchscript for vit-hybrid dynamic_resize --- timm/models/vision_transformer.py | 2 +- timm/models/vision_transformer_hybrid.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/timm/models/vision_transformer.py b/timm/models/vision_transformer.py index b63fe79f07..1de89e4567 100644 --- a/timm/models/vision_transformer.py +++ b/timm/models/vision_transformer.py @@ -38,7 +38,7 @@ from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD, \ OPENAI_CLIP_MEAN, OPENAI_CLIP_STD from timm.layers import PatchEmbed, Mlp, DropPath, trunc_normal_, lecun_normal_, resample_patch_embed, \ - resample_abs_pos_embed, resample_abs_pos_embed_nhwc, RmsNorm, PatchDropout, use_fused_attn, SwiGLUPacked + resample_abs_pos_embed, RmsNorm, PatchDropout, use_fused_attn, SwiGLUPacked from ._builder import build_model_with_cfg from ._manipulate import named_apply, checkpoint_seq, adapt_input_conv from ._registry import generate_default_cfgs, register_model, register_model_deprecations diff --git a/timm/models/vision_transformer_hybrid.py b/timm/models/vision_transformer_hybrid.py index c9b0e7acc7..7f04613e6e 100644 --- a/timm/models/vision_transformer_hybrid.py +++ b/timm/models/vision_transformer_hybrid.py @@ -31,6 +31,8 @@ class HybridEmbed(nn.Module): """ CNN Feature Map Embedding Extract feature map from CNN, flatten, project to embedding dim. """ + output_fmt: Format + def __init__( self, backbone,