Update app.py

app.py

@@ -133,7 +133,8 @@ def encode_and_trace(text, selected_roles):
             masked_input = ids.where(mask_flags, MASK_ID)
 
             encoded_m = encode(masked_input, attn)
-            logits = mlm_head(encoded_m).logits[0]
+            logits = mlm_head(encoded_m)[0]  # shape: (S, V)
+
             preds = logits.argmax(-1)
 
             masked_positions = (~mask_flags[0]).nonzero(as_tuple=False).squeeze(-1)