Minor fix to torch.load() (#4392)

* fix weights only

* change otx version

* change Changelog

* fix linter

Author: kprokofi

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 All notable changes to this project will be documented in this file.
 
+## \[2.4.4\]
+
+### Bug fixes
+
+- Fix torch.load() to be able to load all OTX custom snapshots
+  (<https://github.com/open-edge-platform/training_extensions/pull/4392>)
+
 ## \[2.4.3\]
 
 ### Enhancements

src/otx/__init__.py

@@ -3,7 +3,7 @@
 # Copyright (C) 2024-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "2.4.3"
+__version__ = "2.4.4"
 
 import os
 from pathlib import Path

src/otx/engine/engine.py

@@ -266,7 +266,7 @@ def train(
             # load the model state from the checkpoint incrementally.
             # This means only the model weights are loaded. If there is a mismatch in label_info,
             # perform incremental weight loading for the model's classification layer.
-            ckpt = torch.load(checkpoint)
+            ckpt = torch.load(checkpoint, weights_only=False)
             self.model.load_state_dict_incrementally(ckpt)
 
         with override_metric_callable(model=self.model, new_metric_callable=metric) as model: