From c715c724e7b655319765dcfdf47ef4313b08da37 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Sat, 22 Jun 2024 08:35:30 -0700 Subject: [PATCH] Fix tracing by removing float cast, should end up float anyways --- timm/models/davit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index efcce4f488..09fa9bedfb 100644 --- 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