fix: make CallSequenceTracker unwrap idempotent for bound methods

[FAdrien]

Description

Fixes a TypeError in CallSequenceTracker that breaks total_macs and total_params metrics when the same module’s forward is unwrapped more than once.

Context

ModelArchitectureStats wraps all nn.Module forwards before a single inference run, records the call sequence, then calls unwrap() on each module while computing MACs/params. If a module appears multiple times in the call sequence, unwrap() runs on the same module more than once.

Cause

unwrap() previously restored the module's forward by checking if the it has the __wrapped__ attribute. The first unwrap correctly restores the bound method. As this bound method can still expose the __wrapped__ attribute, a second unwrap then replaces it with the unbound function. Next calls to the model (e.g., to generate an image) then fail with a TypeError because self is no longer bound.

Fix

This PR stores the original forward on a dedicated __call_tracker_original_forward__ attribute, skips re-wrapping modules that are already tracked, and makes unwrap() idempotent on multiple calls.

Related Issue

Fixes #(issue number)

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor (no functional change)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

Testing

• I added or updated tests covering my changes
• Existing tests pass locally (uv run pytest -m "cpu and not slow")

For full setup and testing instructions, see the Contributing Guide.

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my code, especially for agent-assisted changes
• I updated the documentation where necessary

---

Thanks for contributing to Pruna! We're excited to review your work.

New to contributing? Check out our Contributing Guide for everything you need to get started.

Note:

• Draft PRs or PRs without a clear and detailed overview may be delayed.
• Please mark your PR as Ready for Review and ensure the sections above are filled out.
• Contributions that are entirely AI-generated without meaningful human review are discouraged.

[begumcig]

looks super good to me! We can merge once the tests pass :) Thank youuuuuuuu