diff --git a/CHANGELOG.md b/CHANGELOG.md index 7bb8cd98..162600fb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## Unreleased + +### Fixed +- Loading fitted transformer models (`SASRecModel`, `BERT4RecModel`) that were saved on GPU with `load_model`, `load` and `loads` on a machine without GPU ([#327](https://github.com/MTSWebServices/RecTools/pull/327)) + ## [0.19.0] - 11.06.2026 ### Added diff --git a/rectools/models/nn/transformers/base.py b/rectools/models/nn/transformers/base.py index 2183005c..8df88de5 100644 --- a/rectools/models/nn/transformers/base.py +++ b/rectools/models/nn/transformers/base.py @@ -668,7 +668,9 @@ def __getstate__(self) -> object: def __setstate__(self, state: tp.Dict[str, tp.Any]) -> None: if "fitted_checkpoint" in state: - checkpoint = torch.load(io.BytesIO(state["fitted_checkpoint"]), weights_only=False) + # Load to CPU so that a model fitted on GPU can be loaded on a machine without one. + # Weights are moved to the right device before inference and by the trainer when fitting. + checkpoint = torch.load(io.BytesIO(state["fitted_checkpoint"]), map_location="cpu", weights_only=False) loaded = self._model_from_checkpoint(checkpoint) else: loaded = self.from_config(state["model_config"]) diff --git a/tests/models/nn/transformers/test_base.py b/tests/models/nn/transformers/test_base.py index 217a89a7..57cb80cc 100644 --- a/tests/models/nn/transformers/test_base.py +++ b/tests/models/nn/transformers/test_base.py @@ -184,6 +184,26 @@ def test_save_load_for_fitted_model( model.fit(dataset_item_features) assert_save_load_do_not_change_model(model, dataset_item_features) + @pytest.mark.parametrize("model_cls", (SASRecModel, BERT4RecModel)) + def test_load_fitted_model_saved_on_gpu_without_gpu( + self, + model_cls: tp.Type[TransformerModelBase], + dataset: Dataset, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + model = model_cls.from_config({"deterministic": True}) + model.fit(dataset) + + # Save the model as if its tensors were on a GPU + with monkeypatch.context() as m: + m.setattr(torch.serialization, "location_tag", lambda storage: "cuda:0") + data = model.dumps() + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + recovered_model = model_cls.loads(data) + + self._assert_same_reco(model, recovered_model, dataset) + @pytest.mark.parametrize("test_dataset", ("dataset", "dataset_item_features")) @pytest.mark.parametrize("model_cls", (SASRecModel, BERT4RecModel)) @pytest.mark.parametrize(