Skip to content

Use base version check before calling _register_load_state_dict_pre_hook #17030

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Mar 14, 2023
Merged

Use base version check before calling _register_load_state_dict_pre_hook #17030

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c9ace4e
redo
awaelchli Mar 10, 2023
9a26288
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 10, 2023
1d1ac7a
fix import
awaelchli Mar 10, 2023
92333f3
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 10, 2023
e7ccd9b
chlog
awaelchli Mar 10, 2023
7f75c33
Merge remote-tracking branch 'origin/bugfix/base-version-check-2' int…
awaelchli Mar 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/lightning/pytorch/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).

- Fixed an issue where `DistributedSampler.set_epoch` wasn't getting called during `trainer.predict` ([#16785](https://github.com/Lightning-AI/lightning/pull/16785), [#16826](https://github.com/Lightning-AI/lightning/pull/16826))

- Fixed an issue with comparing torch versions when using a version of torch built from source ([#17030](https://github.com/Lightning-AI/lightning/pull/17030))


## [1.9.4] - 2023-03-01

Expand Down
14 changes: 5 additions & 9 deletions src/lightning/pytorch/core/module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""The LightningModule - an nn.Module with many additional features."""

import logging
import numbers
import operator
import weakref
from contextlib import contextmanager
from pathlib import Path
Expand All @@ -37,7 +37,7 @@

import torch
from lightning_utilities.core.apply_func import apply_to_collection
from lightning_utilities.core.imports import RequirementCache
from lightning_utilities.core.imports import compare_version, RequirementCache
from torch import ScriptModule, Tensor
from torch.nn import Module
from torch.optim.optimizer import Optimizer
Expand All @@ -51,12 +51,7 @@
from lightning.fabric.utilities.cloud_io import get_filesystem
from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
from lightning.fabric.utilities.distributed import _distributed_available
from lightning.fabric.utilities.imports import (
_IS_WINDOWS,
_TORCH_GREATER_EQUAL_1_13,
_TORCH_GREATER_EQUAL_2_0,
_TORCH_GREATER_EQUAL_2_1,
)
from lightning.fabric.utilities.imports import _IS_WINDOWS, _TORCH_GREATER_EQUAL_2_0, _TORCH_GREATER_EQUAL_2_1
from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
from lightning.fabric.wrappers import _FabricOptimizer
from lightning.pytorch.callbacks.callback import Callback
Expand Down Expand Up @@ -1567,7 +1562,8 @@ def _register_sharded_tensor_state_dict_hooks_if_available(self) -> None:

self._register_state_dict_hook(state_dict_hook)

if _TORCH_GREATER_EQUAL_1_13:
if compare_version("torch", operator.ge, "1.13.0", use_base_version=True):
# See https://github.com/Lightning-AI/lightning/issues/16644 for why a base-version check is used here
self._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)
else:
# We need to make sure the self inside the method is a weakref proxy
Expand Down