Fix tracing by removing float cast, should end up float anyways

huggingface · Jun 22, 2024 · c715c72 · c715c72
1 parent fb58a73
commit c715c72
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/timm/models/davit.py b/timm/models/davit.py
@@ -137,7 +137,7 @@ def forward(self, x):
         q, k, v = qkv.unbind(0)
 
         if self.dynamic_scale:
-            q = q * float(N) ** -0.5
+            q = q * N ** -0.5
         else:
             q = q * self.head_dim ** -0.5
         attn = q.transpose(-1, -2) @ k