Upload app.py

app.py

@@ -40,12 +40,8 @@ class WaveLmStutterClassification(nn.Module):
                     for param in layer.parameters():
                         param.requires_grad = True
         
-        self.classifier = nn.Sequential(
-            nn.Linear(self.hidden_size, 256),
-            nn.ReLU(),
-            nn.Dropout(0.3),
-            nn.Linear(256, num_labels)
-        )
+        # Single linear layer to match the trained checkpoint
+        self.classifier = nn.Linear(self.hidden_size, num_labels)
         self.num_labels = num_labels
     
     def forward(self, input_values, attention_mask=None):