Remove unsafe torch.load from examples/diffusers/fastgen

[kevalmorabia97]

Follow-up to #1326 - Remove unsafe torch.load(..., weights_only=False) in examples/diffusers/fastgen

Summary by CodeRabbit

• Chores
  • Updated checkpoint loading mechanisms in FastGen examples to improve compatibility and reliability during model restoration.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: cb997879-0eac-4f7d-aa7d-e8bd8256176d

📥 Commits

Reviewing files that changed from the base of the PR and between 55a2101 and 88ed3df.

📒 Files selected for processing (3)

• .github/CODEOWNERS
• examples/diffusers/fastgen/dmd2_recipe.py
• examples/diffusers/fastgen/inference_dmd2_qwen_image.py

---

📝 Walkthrough

Walkthrough

Removes the explicit weights_only=False argument from five torch.load calls in the fastgen example scripts (dmd2_recipe.py and inference_dmd2_qwen_image.py), keeping map_location="cpu" in place. Also adds a CODEOWNERS rule assigning modelopt/torch/fastgen to a new team.

Changes

torch.load weights_only=False Removal in fastgen Examples

Layer / File(s)	Summary
Remove weights_only=False from checkpoint restore calls examples/diffusers/fastgen/dmd2_recipe.py, examples/diffusers/fastgen/inference_dmd2_qwen_image.py	All torch.load calls for EMA shadow, dmd_state.pt, discriminator state, discriminator optimizer state, and EMA inference path drop the weights_only=False argument; map_location="cpu" and surrounding logic remain unchanged.

CODEOWNERS Update

Layer / File(s)	Summary
Add fastgen CODEOWNERS entry .github/CODEOWNERS	New rule maps modelopt/torch/fastgen to the modelopt-torch-fastgen-codeowners team.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Remove unsafe torch.load from examples/diffusers/fastgen' accurately and specifically describes the main change: removing unsafe torch.load calls (specifically those with weights_only=False) from the fastgen example module.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	All security anti-patterns checked per SECURITY.md: no torch.load(..., weights_only=False), numpy.load(allow_pickle=True), trust_remote_code=True, eval/exec on external input, #nosec comments, or n...

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch kmorabia/fix-fastgen-load

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.11%. Comparing base (55a2101) to head (01f0871).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1740      +/-   ##
==========================================
+ Coverage   77.10%   77.11%   +0.01%     
==========================================
  Files         511      511              
  Lines       56260    56260              
==========================================
+ Hits        43380    43387       +7     
+ Misses      12880    12873       -7     

Flag	Coverage Δ
examples	41.84% <ø> (-0.11%)	⬇️
unit	54.39% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 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.

[jingyu-ml]

LGTM

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-06-16 03:14 UTC