mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Make all models FX traceable

Browse Source
This commit is contained in:
Alexander Soare 2021-08-12 15:31:24 +01:00
parent cf4561ca72
commit e051dce354
4 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
timm/models/layers/bottleneck_attn.py
View File

@ -22,6 +22,7 @@ import torch.nn.functional as F
from .helpers import to_2tuple, make_divisible
from .weight_init import trunc_normal_
from timm.models.fx_helpers import fx_and
def rel_logits_1d(q, rel_k, permute_mask: List[int]):

1
timm/models/layers/halo_attn.py
View File

@ -24,7 +24,6 @@ import torch.nn.functional as F
from .helpers import make_divisible
from .weight_init import trunc_normal_
from timm.models.fx_helpers import
def rel_logits_1d(q, rel_k, permute_mask: List[int]):

1
timm/models/layers/non_local_attn.py
View File

@ -10,6 +10,7 @@ from torch.nn import functional as F
from .conv_bn_act import ConvBnAct
from .helpers import make_divisible
from timm.models.fx_helpers import fx_and
class NonLocalAttn(nn.Module):

1
timm/models/tnt.py
View File

@ -12,6 +12,7 @@ import torch.nn as nn
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.helpers import build_model_with_cfg
from timm.models.fx_helpers import fx_and
from timm.models.layers import Mlp, DropPath, trunc_normal_
from timm.models.layers.helpers import to_2tuple
from timm.models.registry import register_model