Skip to content

fix(training): use singular candidate kwargs in test_step - #433

Open
ChrisW09 wants to merge 2 commits into
mainfrom
fix/test-step-candidate-kwargs
Open

fix(training): use singular candidate kwargs in test_step#433
ChrisW09 wants to merge 2 commits into
mainfrom
fix/test-step-candidate-kwargs

Conversation

@ChrisW09

Copy link
Copy Markdown
Collaborator

Part of #413 (the wiring half; the retrieval-leakage half of that issue is not addressed here).

Problem

TaskModel.test_step calls:

preds = self.estimator.predict_with_candidates(
    *data, candidates_x=self.train_features, candidates_y=self.train_targets
)

but both implementations — deeptab/architectures/tabr.py and deeptab/architectures/experimental/modern_nca.py — are declared as def predict_with_candidates(self, *data, candidate_x, candidate_y) with no **kwargs. So trainer.test(...) on any TabR/ModernNCA model always raises TypeError: predict_with_candidates() got an unexpected keyword argument. predict_step (a few lines below) already uses the singular names, which is what the intended spelling is.

Fix

Rename the two kwargs at the call site.

Tests

New tests/test_candidate_models.py drives test_step and predict_step with a minimal stub estimator that has TabR's signature and records what it received. Verified the new test_step test fails on main and passes with the fix.

🤖 Generated with Claude Code

ChrisW09 and others added 2 commits July 27, 2026 21:28
test_step called predict_with_candidates with candidates_x/candidates_y,
but TabR and ModernNCA both define the method as
(self, *data, candidate_x, candidate_y) with no **kwargs, so
trainer.test() on any candidate model raised TypeError. predict_step
already used the correct names.

Part of #413

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
compute_loss is typed as returning float, and estimator is typed
nn.Module; cast for the stub's recorded kwargs.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant