Merge branch 'main' into grad_checkpointing

Author: brianhou0208

tests/test_optim.py

@@ -298,7 +298,7 @@ def test_optim_factory(optimizer):
     assert isinstance(opt_info, OptimInfo)
 
     lr = (1e-2,) * 4
-    if optimizer in ('mars', 'nadam', 'claprop', 'crmsproptf', 'cadafactorbv', 'csgdw', 'clamb'):
+    if optimizer in ('mars', 'nadam', 'claprop', 'crmsproptf', 'cadafactorbv', 'csgdw', 'csgdc', 'clamb'):
         lr = (1e-3,) * 4
     elif optimizer in ('cmars',):
         lr = (1e-4,) * 4
@@ -378,7 +378,7 @@ def test_sgd(optimizer):
     _test_model(optimizer, dict(lr=1e-3))
 
 
-@pytest.mark.parametrize('optimizer',  ['adamw', 'adam', 'nadam', 'adamax', 'nadamw'])
+@pytest.mark.parametrize('optimizer',  ['adamw', 'adam', 'nadam', 'adamax', 'nadamw', 'adamwlegacy', 'adamc'])
 def test_adam(optimizer):
     _test_rosenbrock(
         lambda params: create_optimizer_v2(params, optimizer, lr=5e-2)

timm/layers/attention.py

@@ -154,8 +154,8 @@ def __init__(
             self.k_proj = nn.Linear(dim, attn_dim, bias=qkv_bias)
             self.v_proj = nn.Linear(dim, attn_dim, bias=qkv_bias)
 
-        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
-        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.q_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
         self.attn_drop = nn.Dropout(attn_drop)
         self.norm = norm_layer(attn_dim) if scale_norm else nn.Identity()
         self.proj = nn.Linear(attn_dim, dim)