Update README.md

README.md

@@ -233,7 +233,7 @@ assert set(head_state.keys()) == {"weight", "bias"}
 # Initialize scoring head (single linear layer)
 hidden_size = base_model.config.hidden_size
 scoring_head = torch.nn.Linear(hidden_size, 1)
-scoring_head.load_state_dict(state)
+scoring_head.load_state_dict(head_state)
 
 # Move to device
 device = "cuda" if torch.cuda.is_available() else "cpu"