[tests] Speed up TestWizardModel

[stephenroller]

Patch description
These tests were slow, so just make them do less computation.

Testing steps

Before:

$ pytest tests/nightly/gpu/test_wizard.py -k test_end2end
434.96s call     tests/nightly/gpu/test_wizard.py::TestWizardModel::test_end2end
32.39s setup    tests/nightly/gpu/test_wizard.py::TestWizardModel::test_end2end

$ pytest tests/nightly/gpu/test_wizard -k test_retrieval
39.39s call     tests/nightly/gpu/test_wizard.py::TestWizardModel::test_retrieval
31.89s setup    tests/nightly/gpu/test_wizard.py::TestWizardModel::test_retrieval

After:

40.47s call     tests/nightly/gpu/test_wizard.py::TestWizardModel::test_end2end
36.70s call     tests/nightly/gpu/test_wizard.py::TestWizardModel::test_retrieval