Update codeformer_arch.py

codeformer_arch.py

@@ -75,9 +75,9 @@ def _get_activation_fn(activation):
     raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
 
 class TransformerSALayer(nn.Module):
-    def __init__(self, embed_dim, nhead=8, dim_mlp=2048, dropout=0.0, activation="gelu"):
+    def __init__(self, embed_dim, n_head=8, dim_mlp=2048, dropout=0.0, activation="gelu"):  # Changed nhead to n_head
         super().__init__()
-        self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)
+        self.self_attn = nn.MultiheadAttention(embed_dim, n_head, dropout=dropout)  # Changed nhead to n_head
         self.linear1 = nn.Linear(embed_dim, dim_mlp)
         self.dropout = nn.Dropout(dropout)
         self.linear2 = nn.Linear(dim_mlp, embed_dim)
@@ -159,7 +159,7 @@ class CodeFormer(VQAutoEncoder):
         self.feat_emb = nn.Linear(256, self.dim_embd)
 
         self.ft_layers = nn.Sequential(*[TransformerSALayer(embed_dim=dim_embd, n_head=n_head, dim_mlp=self.dim_mlp, dropout=0.0) 
-                                        for _ in range(self.n_layers)])  # Changed nhead to n_head
+                                        for _ in range(self.n_layers)])
 
         self.idx_pred_layer = nn.Sequential(
             nn.LayerNorm(dim_embd),