torch.fx.wrap not working with older pytorch, trying register_notrace instead

Author: rwightman

timm/layers/attention.py

@@ -8,7 +8,6 @@
 from .pos_embed_sincos import apply_rot_embed_cat
 
 
-@torch.fx.wrap
 def maybe_add_mask(scores: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
     return scores if attn_mask is None else scores + attn_mask
 

timm/models/_features_fx.py

@@ -18,7 +18,7 @@
 
 # Layers we went to treat as leaf modules
 from timm.layers import Conv2dSame, ScaledStdConv2dSame, CondConv2d, StdConv2dSame, Format
-from timm.layers import resample_abs_pos_embed, resample_abs_pos_embed_nhwc
+from timm.layers import resample_abs_pos_embed, resample_abs_pos_embed_nhwc, maybe_add_mask
 from timm.layers.non_local_attn import BilinearAttnTransform
 from timm.layers.pool2d_same import MaxPool2dSame, AvgPool2dSame
 from timm.layers.norm_act import (
@@ -79,6 +79,7 @@ def get_notrace_modules():
 _autowrap_functions = {
     resample_abs_pos_embed,
     resample_abs_pos_embed_nhwc,
+    maybe_add_mask,
 }