tirex-mirror 2025.11.13__py3-none-any.whl → 2025.11.15__py3-none-any.whl

tirex/base.py

@@ -3,6 +3,7 @@
 
 import logging
 import os
+import warnings
 from abc import ABC, abstractmethod
 from typing import Literal, TypeVar
 
@@ -60,7 +61,7 @@ class PretrainedModel(ABC):
         if ckp_kwargs is None:
             ckp_kwargs = {}
         if device is None:
-            device = "cuda:0" if backend == "cuda" else "cpu"
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
         if os.path.exists(path):
             print("Loading weights from local directory")
             checkpoint_path = path
@@ -93,7 +94,7 @@ class PretrainedModel(ABC):
 def load_model(
     path: str,
     device: str | None = None,
-    backend: Literal["torch", "cuda"] | None = None,
+    backend: Literal["torch", "cuda"] = "torch",
     compile: bool = False,
     hf_kwargs=None,
     ckp_kwargs=None,
@@ -115,8 +116,22 @@ def load_model(
         model: ForecastModel = load_model("NX-AI/TiRex")
     """
 
-    if backend is None:
-        backend = "torch" if skip_cuda() or not xlstm_available() else "cuda"
+    if backend == "cuda" and not xlstm_available() and torch.cuda.is_available():
+        backend = "torch"
+        warnings.warn(
+            "Switching to 'torch' backend with device='gpu'. In oder to use a CUDA backend, please make sure that xlstm package is installed.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    if device is not None and device.startswith("cuda") and not torch.cuda.is_available():
+        raise ValueError(
+            "CUDA is not available! This could be because:\n"
+            "  - No GPU is present on this machine\n"
+            "  - GPU drivers are not installed or not functioning\n"
+            "  - PyTorch is installed without CUDA support (CPU-only version)\n"
+            "To resolve: use device='cpu' for CPU inference, or install CUDA-enabled PyTorch if a GPU is available."
+        )
 
     try:
         _, model_string = parse_hf_repo_id(path).split("/")

tirex/models/slstm/cell.py

@@ -70,8 +70,8 @@ class sLSTMCell(nn.Module):
         if input.device.type != "cuda":
             warnings.warn(
                 f"You use TiRex with sLSTM CUDA kernels BUT DO NOT LOAD THE DEVICE ON A CUDA DEVICE (device type is {input.device.type})!"
-                "This is not supported and calls to the model will likely lead to an error if you dont move your model to a CUDA device!"
-                "If you want to run TiRex on CPU you need to disable sLSTM CUDA kernels but be aware of the downsides (see FAQ)"
+                "This is not supported and calls to the model will likely lead to an error if you don't move your model to a CUDA device!"
+                "If you want to run TiRex on CPU, please select backend='torch' and device='cpu'"
             )
 
         if not hasattr(self, "func"):

{tirex_mirror-2025.11.13.dist-info → tirex_mirror-2025.11.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.11.13
+Version: 2025.11.15
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT

{tirex_mirror-2025.11.13.dist-info → tirex_mirror-2025.11.15.dist-info}/RECORD

@@ -1,5 +1,5 @@
 tirex/__init__.py,sha256=rfsOeCJ7eRqU3K3TOhfN5-4XUuZFqt11wBRxk5SoAWA,292
-tirex/base.py,sha256=cwmvqgvqHDFXkVZQ4qczPa6FPsp1c2SWvuByDR2Ecb8,4331
+tirex/base.py,sha256=pGShyI3LwPAwkl0rIYqik5qb8uBAWqbQrVoVjq7UH_8,5096
 tirex/util.py,sha256=ggfETirJ589Dr5o3QThxTnjhgNnMCk10bNoJghnoeoA,31672
 tirex/api_adapter/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
 tirex/api_adapter/forecast.py,sha256=FnRgdI_4vJ7iFTqyWwxPE_C7MBB3nhLTnVMzTr5pihc,17581
@@ -10,12 +10,12 @@ tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
 tirex/models/patcher.py,sha256=8T4c3PZnOAsEpahhrjtt7S7405WUjN6g3cV33E55PD4,1911
 tirex/models/tirex.py,sha256=URt-MClXu0zdUHACQ96Zu3Ytdb52vbeG_SXnj5C4tI8,7522
 tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
-tirex/models/slstm/cell.py,sha256=Q0eRInybFx7Mj_pmlCl3CW9SShT2Qxegrr9KZzy7nrU,7320
+tirex/models/slstm/cell.py,sha256=Otyil_AjpJbUckkINWGHxlqP14J5epm_J_zdWPzvD2g,7290
 tirex/models/slstm/layer.py,sha256=hrDydQJIAHf5W0A0Rt0hXG4yKXrOSY-HPL0UbigR6Q8,2867
-tirex_mirror-2025.11.13.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.11.13.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.11.13.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.11.13.dist-info/METADATA,sha256=QadU-i-QDPdf5Queey80pmF6m96r5NNNBY7N6eGgz-Y,11516
-tirex_mirror-2025.11.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.11.13.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.11.13.dist-info/RECORD,,
+tirex_mirror-2025.11.15.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.11.15.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.11.15.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.11.15.dist-info/METADATA,sha256=eI99-lkrGyhCPfQwFXik7QwZNgsX-qoRR0URpNCdmuM,11516
+tirex_mirror-2025.11.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.11.15.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.11.15.dist-info/RECORD,,