check_version

open-mmlab · Jul 1, 2021 · 2e445ba · 2e445ba
1 parent 043c81d
commit 2e445ba
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/mmaction/models/common/transformer.py b/mmaction/models/common/transformer.py
@@ -1,3 +1,5 @@
+from distutils.version import LooseVersion
+
 import torch
 import torch.nn as nn
 from einops import rearrange
@@ -43,6 +45,9 @@ def __init__(self,
         self.num_heads = num_heads
         self.num_frames = num_frames
         self.norm = build_norm_layer(norm_cfg, self.embed_dims)[1]
+
+        if LooseVersion(torch.__version__) < LooseVersion('1.9.0'):
+            kwargs.pop('batch_first', None)
         self.attn = nn.MultiheadAttention(embed_dims, num_heads, attn_drop,
                                           **kwargs)
         self.proj_drop = nn.Dropout(proj_drop)
@@ -118,6 +123,8 @@ def __init__(self,
         self.num_heads = num_heads
         self.num_frames = num_frames
         self.norm = build_norm_layer(norm_cfg, self.embed_dims)[1]
+        if LooseVersion(torch.__version__) < LooseVersion('1.9.0'):
+            kwargs.pop('batch_first', None)
         self.attn = nn.MultiheadAttention(embed_dims, num_heads, attn_drop,
                                           **kwargs)
         self.proj_drop = nn.Dropout(proj_drop)