Update davit.py

3 years ago · b1337d7f4c
parent 37ab7fe2e2
commit b1337d7f4c
1 changed files with 1 additions and 0 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -565,6 +565,7 @@ def checkpoint_filter_fn(state_dict, model):
        k = re.sub(r'patch_embeds.([0-9]+)', r'stages.\1.patch_embed', k)
        k = re.sub(r'main_blocks.([0-9]+)', r'stages.\1.blocks', k)
        k = k.replace('stages.0.patch_embed', 'stem')
+        k = k.replace('norms.', 'norm.')
        k = k.replace('head.', 'head.fc.')
        k = k.replace('cpe.0', 'cpe1')
        k = k.replace('cpe.1', 'cpe2')