From 3a44e6c60265217122d4bed4e2d46bad2a63ae7a Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Thu, 10 Aug 2023 11:15:58 -0700 Subject: [PATCH] Fix #1912 CoaT model not loading w/ return_interm_layers --- timm/models/coat.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/timm/models/coat.py b/timm/models/coat.py index f6c89af2..68358b3d 100644 --- a/timm/models/coat.py +++ b/timm/models/coat.py @@ -690,8 +690,11 @@ def checkpoint_filter_fn(state_dict, model): for k, v in state_dict.items(): # original model had unused norm layers, removing them requires filtering pretrained checkpoints if k.startswith('norm1') or \ - (model.norm2 is None and k.startswith('norm2')) or \ - (model.norm3 is None and k.startswith('norm3')): + (k.startswith('norm2') and getattr(model, 'norm2', None) is None) or \ + (k.startswith('norm3') and getattr(model, 'norm3', None) is None) or \ + (k.startswith('norm4') and getattr(model, 'norm4', None) is None) or \ + (k.startswith('aggregate') and getattr(model, 'aggregate', None) is None) or \ + (k.startswith('head') and getattr(model, 'head', None) is None): continue out_dict[k] = v return out_dict