Fix ModelPruning iterative reuse RuntimeError on non-leaf tensors

[Priyanshu-byte-coder]

What does this PR fix?

Fixes #8542.

When ModelPruning(use_lottery_ticket_hypothesis=True) is reused across multiple trainer.fit() calls (iterative pruning), the second and subsequent calls to setup() raise:

RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment

Root cause: After the first training pass, apply_lottery_ticket_hypothesis() calls _copy_param() which does dst.data = src.data.to(dst.device). This in-place data assignment makes the model parameters non-leaf tensors. When setup() is called again on the next trainer.fit(), deepcopy(module) fails.

Changes

src/lightning/pytorch/callbacks/pruning.py

• Added _deepcopy_for_pruning(module) static helper: temporarily replaces non-leaf parameters with detach().clone() copies, deepcopies, then restores originals
• Updated setup() to use _deepcopy_for_pruning instead of bare deepcopy(module)
• In setup(), set self._original_layers = None before re-populating to release previous-run tensor references

tests/tests_pytorch/callbacks/test_pruning.py

• Added test_iterative_pruning_no_runtime_error: runs 3 consecutive trainer.fit() calls with the same pruning callback and verifies no RuntimeError is raised

Reproduction

model = BoringModel()
pruning_callback = ModelPruning("l1_unstructured", use_lottery_ticket_hypothesis=True, amount=0.2)
for _ in range(2):
    trainer = Trainer(max_epochs=1, accelerator="cpu", callbacks=[pruning_callback])
    trainer.fit(model)  # RuntimeError on second iteration (before fix)

---

📚 Documentation preview 📚: https://pytorch-lightning--21717.org.readthedocs.build/en/21717/

[Priyanshu-byte-coder]

Pinging for review — all CI checks are green (triage, pre-commit, GitGuardian, ReadTheDocs all pass).

@justusschock @tchaton could one of you take a look when you get a chance? The fix is in _deepcopy_for_pruning in pruning.py — patches module.__dict__ instead of module._parameters to handle the non-leaf tensor that PyTorch's pruning forward pre-hook stores there.

[Priyanshu-byte-coder]

Follow-up ping @justusschock @tchaton — CI is green and the fix is small: ModelPruning no longer raises RuntimeError on non-leaf tensors during iterative reuse (pruning.py + regression test). Happy to rebase onto current master if that helps get it reviewed.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79%. Comparing base (4819088) to head (b528ac5).
✅ All tests successful. No failed tests found.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

❗ There is a different number of reports uploaded between BASE (4819088) and HEAD (b528ac5). Click for more details.

HEAD has 1666 uploads less than BASE

Flag	BASE (4819088)	HEAD (b528ac5)
python3.10	30	3
cpu	417	42
lightning	149	15
pytest	209	0
python	30	3
lightning_fabric	134	0
python3.12	120	12
python3.12.7	88	9
python3.11	60	6
python3.13	89	9
pytorch2.8	30	6
pytorch_lightning	134	27
pytest-full	208	42
pytorch2.7	15	3
pytorch2.9	30	6
pytorch2.10	30	6
pytorch2.6	14	3
pytorch2.4.1	14	3
pytorch2.2.2	15	3
pytorch2.1	30	6
pytorch2.3	15	3
pytorch2.5.1	15	3

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #21717     +/-   ##
=========================================
- Coverage      87%      79%     -8%     
=========================================
  Files         270      267      -3     
  Lines       24010    23963     -47     
=========================================
- Hits        20790    18842   -1948     
- Misses       3220     5121   +1901