Required prerequisites

• I have read the documentation https://torchopt.readthedocs.io.
• I have searched the Issue Tracker and Discussions that this hasn't already been reported. (+1 or comment there if it has.)
• Consider asking first in a Discussion.

What version of TorchOpt are you using?

0.7.0

System information

Problem description

torchopt.utils.extract_state_dict has an incorrect overloading annotation. The detach_buffers argument is missing from the overloading stubs, meaning that the static type checker complains when using detach_buffers.

Reproducible example code

The Python snippets:

torchopt.extract_state_dict(
    nn.Linear(1,1), by="reference", detach_buffers=True
)

Traceback

No response

Expected behavior

Type checker correctly checks the intended usage of extract_state_dict.

Additional context

No response