Load pretrained weights notebook and bug

[drewoldag]

Found a bug that prevented continued training with pretrained weights.
Added initial notebook, fixed bug, updated test.

Change Description

Closes #742

Solution Description

Code Quality

• I have read the Contribution Guide and agree to the Code of Conduct
• My code follows the code style of this project
• My code builds (or compiles) cleanly without any errors or warnings
• My code contains relevant comments and necessary documentation

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[github-actions]

Before [f0869f2]	After [e701b0a]	Ratio	Benchmark (Parameter)
failed	failed	n/a	data_cache_benchmarks.DataCacheBenchmarks.time_preload_cache_hsc1k
failed	failed	n/a	data_cache_benchmarks.DataCacheBenchmarks.track_cache_hsc1k_hyrax_size_undercount
failed	failed	n/a	data_request_benchmarks.DatasetRequestBenchmarks.time_request_all_data
38.3±0.1ms	39.5±0.6ms	1.03	benchmarks.time_nb_obj_construct
39.2±0.3ms	39.8±0.6ms	1.02	benchmarks.time_nb_obj_dir
3.74G	3.83G	1.02	vector_db_benchmarks.VectorDBInsertBenchmarks.peakmem_load_vector_db(16384, 'qdrant')
10.00±0.04ms	10.2±0.03ms	1.02	vector_db_benchmarks.VectorDBSearchBenchmarks.time_search_by_vector_many_shards(64, 'chromadb')
1.93±0.01s	1.95±0.01s	1.01	benchmarks.time_help
1.96±0.01s	1.98±0.01s	1.01	benchmarks.time_prepare_help
1.56G	1.57G	1.01	vector_db_benchmarks.VectorDBInsertBenchmarks.peakmem_load_vector_db(16384, 'chromadb')

Click here to view all benchmarks.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 64.69%. Comparing base (f0869f2) to head (cea1023).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #770   +/-   ##
=======================================
  Coverage   64.69%   64.69%           
=======================================
  Files          61       61           
  Lines        5894     5894           
=======================================
  Hits         3813     3813           
  Misses       2081     2081           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Copilot]

Pull request overview

This PR fixes a bug that prevented continued training with pretrained weights and adds a demonstration notebook for loading pretrained model weights.

Changes:

• Bug fix: Removed the assign=True parameter from load_state_dict call in model_registry.py
• Test update: Added epoch config update to test_train_resume to properly cover the resumed training scenario
• New notebook: Added load_pretrained_model.ipynb demonstrating fine-tuning from saved weights, with a corresponding entry in common_workflows.rst

Reviewed changes

Copilot reviewed 4 out of 5 changed files in this pull request and generated 1 comment.

File	Description
src/hyrax/models/model_registry.py	Core bug fix — removes assign=True from load_state_dict
tests/hyrax/test_train.py	Updates test to set epoch count before resuming
docs/pre_executed/load_pretrained_model.ipynb	New notebook demonstrating pretrained weight loading workflow
docs/common_workflows.rst	Registers the new notebook in the documentation table of contents

---

You can also share your feedback on Copilot code review. Take the survey.

[drewoldag]

The fact that sphinx re-ran the notebook (despite being in the pre-executed directory) causes the overly verbose look of the notebook in the docs.

I'll try to track down what is causing sphinx to run the notebook when it shouldn't - for now let's block the merge.

Once the notebook looks bite-sized again, we can move ahead.

[aritraghsh09]

One comment. Not blocking. But please take a look before merging.

[aritraghsh09]

I just want to check that this won't cause problems when someone is trying to resume from a checkpoint (as opposed to running from a weights file)? If we haven't tested it yet, maybe we should test it on this branch?

[drewoldag]

I double checked and was able to reproduce the "hyrax checkpointing" notebook. So I don't think we'll be introducing a regression with this change.

[drewoldag]

Additionally, the code that resumes from a checkpoint doesn't use this codepath, so we should be double ok :)