Update davit.py

3 years ago · 28b712cf5b
parent f32c511c6a
commit 28b712cf5b
1 changed files with 1 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -341,7 +341,7 @@ class SpatialBlock(nn.Module):
        x = self.cpe[1](x, size)
        if self.ffn:
            x = x + self.drop_path(self.mlp(self.norm2(x)))
-        return x : Tensor, size
+        return x, size