Weird behaviour, need to repeat wrap?

Author: rwightman

timm/layers/attention_pool2d.py

@@ -12,12 +12,16 @@
 import torch
 import torch.nn as nn
 
-from. config import use_fused_attn
+from ._fx import register_notrace_function
+from .config import use_fused_attn
 from .helpers import to_2tuple
 from .pos_embed import resample_abs_pos_embed
 from .pos_embed_sincos import apply_rot_embed, RotaryEmbedding
 from .weight_init import trunc_normal_
 
+# have to register again for some reason
+register_notrace_function(resample_abs_pos_embed)
+
 
 class RotAttentionPool2d(nn.Module):
     """ Attention based 2D feature pooling w/ rotary (relative) pos embedding.