Explicitly set weights_only=False for torch version 2.4+

Author: opcode81

src/sensai/torch/torch_base.py

@@ -3,6 +3,7 @@
 import logging
 import typing
 from abc import ABC, abstractmethod
+from platform import version
 from typing import Union, Tuple, Callable, Optional, List, Sequence
 
 import numpy as np
@@ -21,9 +22,11 @@
 from ..util.dtype import to_float_array
 from ..util.pickle import setstate
 from ..util.string import ToStringMixin
+from ..util.version import Version
 from ..vector_model import VectorRegressionModel, VectorClassificationModel, TrainingContext
 
 log: logging.Logger = logging.getLogger(__name__)
+torch_version = Version(torch)
 
 
 class MCDropoutCapableNNModule(nn.Module, ABC):
@@ -139,9 +142,12 @@ def _set_cuda_enabled(self, is_cuda_enabled: bool) -> None:
     def _is_cuda_enabled(self) -> bool:
         return self.cuda
 
-    def _load_model(self, model_file) -> None:  # TODO: complete type hints: what types are allowed for modelFile?
+    def _load_model(self, model_file) -> None:
         try:
-            self.module = torch.load(model_file)
+            load_kwargs = {}
+            if torch_version.is_at_least(2, 4):
+                load_kwargs["weights_only"] = False
+            self.module = torch.load(model_file, **load_kwargs)
             self._gpu = self._get_gpu_from_model_parameter_device()
         except:
             if self._is_cuda_enabled():