Changing _modules dict type in Pytorch 2.5 leads to failing collections test

[bfolie]

🐛 Bug

As of PyTorch 2.5, the Module._modules attribute has changed from an OrderedDict to a dict. This was done to improve the performance of certain TorchDynamo guard mechanisms and is not expected to affect other behavior because, as of Python 3.7, dictionaries are ordered by default.

This means that MetricCollection.keys returns type dict_keys when keep_base=True but returns odict_keys when keep_base=False because it uses _to_renamed_ordered_dict to construct a new OrderedDict (MetricCollection.items has the same issue).

As a result, test_metric_collection_prefix_postfix_args fails at line 200 because it calls .keys() with keep_base equal to both True and False and asserts that the resulting types are the same, which they are not.

Proposed Fix

Change _to_renamed_ordered_dict so that it produces either a dict or an OrderedDict, depending on the type of the existing _modules dictionary. This will work for both older and newer version of PyTorch.

Current:

def _to_renamed_ordered_dict(self) -> OrderedDict:
    od = OrderedDict()
    for k, v in self._modules.items():
        od[self._set_name(k)] = v
    return od

Proposed:

def _to_renamed_dict(self) -> dict:
    if isinstance(self._modules, OrderedDict):
        d = OrderedDict()
    else:
        d = dict()
    for k, v in self._modules.items():
        d[self._set_name(k)] = v
    return d

To Reproduce

Run the unit tests with PyTorch 2.5

[github-actions]

Hi! thanks for your contribution!, great first issue!

[bfolie]

I am happy to implement the proposed fix

[SkafteNicki]

Hi @bfolie, we be happy to receive a fix from you :)

[Borda]

@bfolie could you pls send a PR with this fix asap so we can rollout the compatibility check for PT 2.5

[bfolie]

Hi @Borda , @SkafteNicki . I made the proposed change and confirmed locally that the test passes using both PT 2.4.1 and 2.5.0. But when I try to push my branch I get an error

remote: Permission to Lightning-AI/torchmetrics.git denied to bfolie.
fatal: unable to access 'https://github.com/Lightning-AI/torchmetrics/': The requested URL returned error: 403

Are there restrictions on who can create new branches?

[SkafteNicki]

Hi @Borda , @SkafteNicki . I made the proposed change and confirmed locally that the test passes using both PT 2.4.1 and 2.5.0. But when I try to push my branch I get an error

remote: Permission to Lightning-AI/torchmetrics.git denied to bfolie.
fatal: unable to access 'https://github.com/Lightning-AI/torchmetrics/': The requested URL returned error: 403

Are there restrictions on who can create new branches?

Yes only members of lightning organizations can push directly. You need to create a fork of the repository, and then push your changes to that. Then you should be able to create a PR from your fork against the main repo