Fix: LlamaRMSNorm to accept extra kwargs from timm

models/ACMDM.py

@@ -424,7 +424,8 @@ class TimestepEmbedder(nn.Module):
 
 
 class LlamaRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
+    def __init__(self, hidden_size, eps=1e-6, **kwargs):
+        # Accept and ignore extra kwargs (like 'device') that timm may pass
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps