Final architecture alignment: BatchNorm(64) for temporal normalization

model.py

@@ -6,13 +6,13 @@ class SignVLM(nn.Module):
     def __init__(self, input_dim=225, hidden_dim=512, num_heads=8, num_layers=4, num_classes=60):
         super(SignVLM, self).__init__()
         
-        # Updated hidden_dim to 512 to match the checkpoint
+        # Layer 0: Projection to 512
         self.feature_extractor = nn.Sequential(
             nn.Linear(input_dim, hidden_dim),
-            nn.BatchNorm1d(hidden_dim)
+            # Layer 1: BatchNorm across the 64 frames
+            nn.BatchNorm1d(64) 
         )
         
-        # Updated dim_feedforward to 1024 (hidden_dim * 2)
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=hidden_dim, 
             nhead=num_heads, 
@@ -22,17 +22,17 @@ class SignVLM(nn.Module):
         self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
         
         self.classifier = nn.Sequential(
-            nn.Linear(hidden_dim, 256), # First layer reduces 512 -> 256
+            nn.Linear(hidden_dim, 256),
             nn.ReLU(),
             nn.Dropout(0.5),
             nn.Linear(256, num_classes)
         )
 
     def forward(self, x):
-        batch_size, seq_len, features = x.shape
-        x = x.view(-1, features)
-        x = self.feature_extractor(x)
-        x = x.view(batch_size, seq_len, -1)
+        # x: [batch, 64, 225]
+        x = self.feature_extractor[0](x) # Linear projection: [batch, 64, 512]
+        x = self.feature_extractor[1](x) # BatchNorm1d: [batch, 64, 512] (normalizing dim 1)
+        
         x = self.transformer(x)
-        x = x.mean(dim=1)
+        x = x.mean(dim=1) # Global Average Pooling
         return self.classifier(x)