tirex-mirror 2025.11.29__py3-none-any.whl → 2025.12.16__py3-none-any.whl

tirex/models/base/base_classifier.py

@@ -0,0 +1,36 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+
+from .base_tirex import BaseTirexEmbeddingModel
+
+
+class BaseTirexClassifier(BaseTirexEmbeddingModel):
+    """Abstract base class for TiRex classification models."""
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class labels for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted class labels with shape (batch_size,).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict(emb)).long()
+
+    @torch.inference_mode()
+    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class probabilities for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict_proba(emb))

tirex/models/base/base_regressor.py

@@ -0,0 +1,23 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+
+from .base_tirex import BaseTirexEmbeddingModel
+
+
+class BaseTirexRegressor(BaseTirexEmbeddingModel):
+    """Abstract base class for TiRex regression models."""
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict values for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted values with shape (batch_size,).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict(emb)).float()

tirex/models/base/base_tirex.py

@@ -0,0 +1,109 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch
+
+from tirex.util import train_val_split
+
+from ..embedding import TiRexEmbedding
+
+
+class BaseTirexEmbeddingModel(ABC):
+    """Abstract base class for TiRex models.
+
+    This base class provides common functionality for all TiRex classifier and regression models,
+    including embedding model initialization and a consistent interface.
+
+    """
+
+    def __init__(
+        self, data_augmentation: bool = False, device: str | None = None, compile: bool = False, batch_size: int = 512
+    ) -> None:
+        """Initializes a base TiRex model.
+
+        This base class initializes the embedding model and common configuration
+        used by both classification and regression models.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+        """
+
+        # Set device
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.device = device
+        self._compile = compile
+
+        self.batch_size = batch_size
+        self.data_augmentation = data_augmentation
+        self.emb_model = TiRexEmbedding(
+            device=self.device,
+            data_augmentation=self.data_augmentation,
+            batch_size=self.batch_size,
+            compile=self._compile,
+        )
+
+    @abstractmethod
+    def fit(self, train_data: tuple[torch.Tensor, torch.Tensor]) -> None:
+        """Abstract method for model training"""
+        pass
+
+    def _compute_embeddings(self, x: torch.Tensor) -> np.ndarray:
+        """Compute embeddings for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+
+        Returns:
+            np.ndarray: Embeddings with shape (batch_size, embedding_dim).
+        """
+        self.emb_model.eval()
+        x = x.to(self.device)
+        return self.emb_model(x).cpu().numpy()
+
+    def _create_train_val_datasets(
+        self,
+        train_data: tuple[torch.Tensor, torch.Tensor],
+        val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
+        val_split_ratio: float = 0.2,
+        stratify: bool = False,
+        seed: int | None = None,
+    ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+        if val_data is None:
+            train_data, val_data = train_val_split(
+                train_data=train_data, val_split_ratio=val_split_ratio, stratify=stratify, seed=seed
+            )
+        return train_data, val_data
+
+    @abstractmethod
+    def save_model(self, path: str) -> None:
+        """Save model to file.
+
+        Args:
+            path: File path where the model should be saved.
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    def load_model(cls, path: str):
+        """Load model from file.
+
+        Args:
+            path: File path to the saved model checkpoint.
+
+        Returns:
+            Instance of the model class with loaded weights and configuration.
+        """
+        pass

tirex/models/classification/__init__.py

@@ -1,8 +1,8 @@
-# classification/__init__.py
-from .linear_classifier import TirexClassifierTorch
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from .gbm_classifier import TirexGBMClassifier
+from .linear_classifier import TirexLinearClassifier
 from .rf_classifier import TirexRFClassifier
 
-__all__ = [
-    "TirexClassifierTorch",
-    "TirexRFClassifier",
-]
+__all__ = ["TirexLinearClassifier", "TirexRFClassifier", "TirexGBMClassifier"]

tirex/models/classification/gbm_classifier.py

@@ -0,0 +1,188 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import joblib
+import torch
+from lightgbm import LGBMClassifier, early_stopping
+
+from ..base.base_classifier import BaseTirexClassifier
+
+
+class TirexGBMClassifier(BaseTirexClassifier):
+    """
+    A Gradient Boosting classifier that uses time series embeddings as features.
+
+    This classifier combines a pre-trained embedding model for feature extraction with a
+    Gradient Boosting classifier.
+
+    Example:
+        >>> from tirex.models.classification import TirexGBMClassifier
+        >>>
+        >>> # Create model with custom LightGBM parameters
+        >>> model = TirexGBMClassifier(
+        ...     data_augmentation=True,
+        ...     n_estimators=50,
+        ...     random_state=42
+        ... )
+        >>>
+        >>> # Prepare data (can use NumPy arrays or PyTorch tensors)
+        >>> X_train = torch.randn(100, 1, 128)  # 100 samples, 1 number of variates, 128 sequence length
+        >>> y_train = torch.randint(0, 3, (100,))  # 3 classes
+        >>>
+        >>> # Train the model
+        >>> model.fit((X_train, y_train))
+        >>>
+        >>> # Make predictions
+        >>> X_test = torch.randn(20, 1, 128)
+        >>> predictions = model.predict(X_test)
+        >>> probabilities = model.predict_proba(X_test)
+    """
+
+    def __init__(
+        self,
+        data_augmentation: bool = False,
+        device: str | None = None,
+        compile: bool = False,
+        batch_size: int = 512,
+        early_stopping_rounds: int | None = 10,
+        min_delta: float = 0.0,
+        val_split_ratio: float = 0.2,
+        stratify: bool = True,
+        # LightGBM kwargs
+        **lgbm_kwargs,
+    ) -> None:
+        """Initializes Embedding Based Gradient Boosting Classification model.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+            early_stopping_rounds: int | None
+                Number of rounds without improvement of all metrics for Early Stopping. Default: 10
+            min_delta: float
+                Minimum improvement in score to keep training. Default 0.0
+            val_split_ratio : float
+                Proportion of training data to use for validation, if validation data are not provided. Default: 0.2
+            stratify : bool
+                Whether to stratify the train/validation split by class labels. Default: True
+            **lgbm_kwargs
+                Additional keyword arguments to pass to LightGBM's LGBMClassifier.
+                Common options include n_estimators, max_depth, learning_rate, random_state, etc.
+        """
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
+
+        # Early Stopping callback
+        self.early_stopping_rounds = early_stopping_rounds
+        self.min_delta = min_delta
+
+        # Data split parameters:
+        self.val_split_ratio = val_split_ratio
+        self.stratify = stratify
+
+        # Extract random_state for train_val_split if provided
+        self.random_state = lgbm_kwargs.get("random_state", None)
+
+        self.head = LGBMClassifier(**lgbm_kwargs)
+
+    @torch.inference_mode()
+    def fit(
+        self,
+        train_data: tuple[torch.Tensor, torch.Tensor],
+        val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> None:
+        """Train the LightGBM classifier on embedded time series data.
+
+        This method generates embeddings for the training data using the embedding
+        model, then trains the LightGBM classifier on these embeddings.
+
+        Args:
+            train_data: Tuple of (X_train, y_train) where X_train is the input time
+                series data (torch.Tensor) and y_train is a torch.Tensor
+                of class labels.
+            val_data: Optional tuple of (X_val, y_val) for validation where X_train is the input time
+                series data (torch.Tensor) and y_train is a torch.Tensor
+                of class labels. If None, validation data will be automatically split from train_data (20% split).
+        """
+
+        (X_train, y_train), (X_val, y_val) = self._create_train_val_datasets(
+            train_data=train_data,
+            val_data=val_data,
+            val_split_ratio=self.val_split_ratio,
+            stratify=self.stratify,
+            seed=self.random_state,
+        )
+
+        X_train = X_train.to(self.device)
+        X_val = X_val.to(self.device)
+
+        embeddings_train = self._compute_embeddings(X_train)
+        embeddings_val = self._compute_embeddings(X_val)
+
+        y_train = y_train.detach().cpu().numpy() if isinstance(y_train, torch.Tensor) else y_train
+        y_val = y_val.detach().cpu().numpy() if isinstance(y_val, torch.Tensor) else y_val
+
+        self.head.fit(
+            embeddings_train,
+            y_train,
+            eval_set=[(embeddings_val, y_val)],
+            callbacks=[early_stopping(stopping_rounds=self.early_stopping_rounds, min_delta=self.min_delta)]
+            if self.early_stopping_rounds is not None
+            else None,
+        )
+
+    def save_model(self, path: str) -> None:
+        """This method saves the trained LightGBM classifier head (joblib format) and embedding information.
+
+        Args:
+            path: File path where the model should be saved (e.g., 'model.joblib').
+        """
+
+        payload = {
+            "data_augmentation": self.data_augmentation,
+            "compile": self._compile,
+            "batch_size": self.batch_size,
+            "early_stopping_rounds": self.early_stopping_rounds,
+            "min_delta": self.min_delta,
+            "val_split_ratio": self.val_split_ratio,
+            "stratify": self.stratify,
+            "head": self.head,
+        }
+        joblib.dump(payload, path)
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexGBMClassifier":
+        """Load a saved model from file.
+
+        This reconstructs the model with the embedding configuration and loads
+        the trained LightGBM classifier from a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexGBMClassifier: The loaded model with trained Gradient Boosting, ready for inference.
+        """
+        checkpoint = joblib.load(path)
+
+        # Create new instance with saved configuration
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
+            batch_size=checkpoint["batch_size"],
+            early_stopping_rounds=checkpoint["early_stopping_rounds"],
+            min_delta=checkpoint["min_delta"],
+            val_split_ratio=checkpoint["val_split_ratio"],
+            stratify=checkpoint["stratify"],
+        )
+
+        # Load the trained LightGBM head
+        model.head = checkpoint["head"]
+
+        # Extract random_state from the loaded head if available
+        model.random_state = getattr(model.head, "random_state", None)
+
+        return model

tirex/models/classification/linear_classifier.py

@@ -1,12 +1,15 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
 from dataclasses import asdict
 
 import torch
 
-from .embedding import TiRexEmbedding
-from .trainer import TrainConfig, Trainer
+from ..base.base_classifier import BaseTirexClassifier
+from ..trainer import TrainConfig, Trainer, TrainingMetrics
 
 
-class TirexClassifierTorch(torch.nn.Module):
+class TirexLinearClassifier(BaseTirexClassifier, torch.nn.Module):
     """
     A PyTorch classifier that combines time series embeddings with a linear classification head.
 
@@ -16,10 +19,10 @@ class TirexClassifierTorch(torch.nn.Module):
 
     Example:
         >>> import torch
-        >>> from tirex.models.classification import TirexClassifierTorch
+        >>> from tirex.models.classification import TirexLinearClassifier
         >>>
-        >>> # Create model with TIREX embeddings
-        >>> model = TirexClassifierTorch(
+        >>> # Create model with TiRex embeddings
+        >>> model = TirexLinearClassifier(
         ...     data_augmentation=True,
         ...     max_epochs=2,
         ...     lr=1e-4,
@@ -43,8 +46,9 @@ class TirexClassifierTorch(torch.nn.Module):
         self,
         data_augmentation: bool = False,
         device: str | None = None,
+        compile: bool = False,
         # Training parameters
-        max_epochs: int = 50,
+        max_epochs: int = 10,
         lr: float = 1e-4,
         weight_decay: float = 0.01,
         batch_size: int = 512,
@@ -62,11 +66,13 @@ class TirexClassifierTorch(torch.nn.Module):
 
         Args:
             data_augmentation : bool | None
-                Whether to use data_augmentation for embeddings (stats and first-order differences of the original data). Default: False
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
             device : str | None
-                Device to run the model on. If None, uses CUDA if available, else CPU. Default: None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
             max_epochs : int
-                Maximum number of training epochs. Default: 50
+                Maximum number of training epochs. Default: 10
             lr : float
                 Learning rate for the optimizer. Default: 1e-4
             weight_decay : float
@@ -91,15 +97,9 @@ class TirexClassifierTorch(torch.nn.Module):
                 Dropout probability for the classification head. If None, no dropout is used. Default: None
         """
 
-        super().__init__()
-
-        if device is None:
-            device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        self.device = device
+        torch.nn.Module.__init__(self)
 
-        # Create embedding model
-        self.emb_model = TiRexEmbedding(device=self.device, data_augmentation=data_augmentation, batch_size=batch_size)
-        self.data_augmentation = data_augmentation
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
 
         # Head parameters
         self.dropout = dropout
@@ -115,6 +115,7 @@ class TirexClassifierTorch(torch.nn.Module):
             lr=lr,
             weight_decay=weight_decay,
             class_weights=class_weights,
+            task_type="classification",
             batch_size=batch_size,
             val_split_ratio=val_split_ratio,
             stratify=stratify,
@@ -132,9 +133,7 @@ class TirexClassifierTorch(torch.nn.Module):
 
     @torch.inference_mode()
     def _identify_head_dims(self, x: torch.Tensor, y: torch.Tensor) -> None:
-        self.emb_model.eval()
-        sample_emb = self.emb_model(x[:1])
-        self.emb_dim = sample_emb.shape[-1]
+        self.emb_dim = self._compute_embeddings(x[:1]).shape[-1]
         self.num_classes = len(torch.unique(y))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -155,7 +154,7 @@ class TirexClassifierTorch(torch.nn.Module):
 
     def fit(
         self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
-    ) -> dict[str, float]:
+    ) -> TrainingMetrics:
         """Train the classification head on the provided data.
 
         This method initializes the classification head based on the data dimensions,
@@ -221,6 +220,7 @@ class TirexClassifierTorch(torch.nn.Module):
             {
                 "head_state_dict": self.head.state_dict(),  # need to save only head, embedding is frozen
                 "data_augmentation": self.data_augmentation,
+                "compile": self._compile,
                 "emb_dim": self.emb_dim,
                 "num_classes": self.num_classes,
                 "dropout": self.dropout,
@@ -230,7 +230,7 @@ class TirexClassifierTorch(torch.nn.Module):
         )
 
     @classmethod
-    def load_model(cls, path: str) -> "TirexClassifierTorch":
+    def load_model(cls, path: str) -> "TirexLinearClassifier":
         """Load a saved model from file.
 
         This reconstructs the model architecture and loads the trained weights from
@@ -239,7 +239,7 @@ class TirexClassifierTorch(torch.nn.Module):
         Args:
             path: File path to the saved model checkpoint.
         Returns:
-            TirexClassifierTorch: The loaded model with trained weights, ready for inference.
+            TirexLinearClassifier: The loaded model with trained weights, ready for inference.
         """
         checkpoint = torch.load(path)
 
@@ -248,6 +248,7 @@ class TirexClassifierTorch(torch.nn.Module):
 
         model = cls(
             data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
             dropout=checkpoint["dropout"],
             max_epochs=train_config_dict.get("max_epochs", 50),
             lr=train_config_dict.get("lr", 1e-4),

tirex/models/classification/rf_classifier.py

@@ -1,12 +1,14 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
 import joblib
-import numpy as np
 import torch
 from sklearn.ensemble import RandomForestClassifier
 
-from .embedding import TiRexEmbedding
+from ..base.base_classifier import BaseTirexClassifier
 
 
-class TirexRFClassifier:
+class TirexRFClassifier(BaseTirexClassifier):
     """
     A Random Forest classifier that uses time series embeddings as features.
 
@@ -15,7 +17,6 @@ class TirexRFClassifier:
     time series, which are then used to train the Random Forest.
 
     Example:
-        >>> import numpy as np
         >>> from tirex.models.classification import TirexRFClassifier
         >>>
         >>> # Create model with custom Random Forest parameters
@@ -43,6 +44,7 @@ class TirexRFClassifier:
         self,
         data_augmentation: bool = False,
         device: str | None = None,
+        compile: bool = False,
         batch_size: int = 512,
         # Random Forest parameters
         **rf_kwargs,
@@ -51,24 +53,18 @@ class TirexRFClassifier:
 
         Args:
             data_augmentation : bool
-                Whether to use data_augmentation for embeddings (stats and first-order differences of the original data). Default: False
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
             device : str | None
                 Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
             batch_size : int
                 Batch size for embedding calculations. Default: 512
             **rf_kwargs
                 Additional keyword arguments to pass to sklearn's RandomForestClassifier.
                 Common options include n_estimators, max_depth, min_samples_split, random_state, etc.
         """
-
-        # Set device
-        if device is None:
-            device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        self.device = device
-
-        self.emb_model = TiRexEmbedding(device=self.device, data_augmentation=data_augmentation, batch_size=batch_size)
-        self.data_augmentation = data_augmentation
-
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
         self.head = RandomForestClassifier(**rf_kwargs)
 
     @torch.inference_mode()
@@ -80,7 +76,7 @@ class TirexRFClassifier:
 
         Args:
             train_data: Tuple of (X_train, y_train) where X_train is the input time
-                series data (torch.Tensor) and y_train is a numpy array
+                series data (torch.Tensor) and y_train is a torch.Tensor
                 of class labels.
         """
         X_train, y_train = train_data
@@ -88,35 +84,10 @@ class TirexRFClassifier:
         if isinstance(y_train, torch.Tensor):
             y_train = y_train.detach().cpu().numpy()
 
-        embeddings = self.emb_model(X_train).cpu().numpy()
-        self.head.fit(embeddings, y_train)
-
-    @torch.inference_mode()
-    def predict(self, x: torch.Tensor) -> torch.Tensor:
-        """Predict class labels for input time series data.
-
-        Args:
-            x: Input time series data as torch.Tensor or np.ndarray with shape
-                (batch_size, num_variates, seq_len).
-        Returns:
-            torch.Tensor: Predicted class labels with shape (batch_size,).
-        """
-
-        embeddings = self.emb_model(x).cpu().numpy()
-        return torch.from_numpy(self.head.predict(embeddings)).long()
-
-    @torch.inference_mode()
-    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
-        """Predict class probabilities for input time series data.
+        X_train = X_train.to(self.device)
+        embeddings = self._compute_embeddings(X_train)
 
-        Args:
-            x: Input time series data as torch.Tensor or np.ndarray with shape
-                (batch_size, num_variates, seq_len).
-        Returns:
-            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
-        """
-        embeddings = self.emb_model(x).cpu().numpy()
-        return torch.from_numpy(self.head.predict_proba(embeddings))
+        self.head.fit(embeddings, y_train)
 
     def save_model(self, path: str) -> None:
         """This method saves the trained Random Forest classifier head and embedding information in joblib format
@@ -126,6 +97,8 @@ class TirexRFClassifier:
         """
         payload = {
             "data_augmentation": self.data_augmentation,
+            "compile": self._compile,
+            "batch_size": self.batch_size,
             "head": self.head,
         }
         joblib.dump(payload, path)
@@ -147,6 +120,8 @@ class TirexRFClassifier:
         # Create new instance with saved configuration
         model = cls(
             data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
+            batch_size=checkpoint["batch_size"],
         )
 
         # Load the trained Random Forest head

tirex/models/{classification/embedding.py → embedding.py}

@@ -1,27 +1,29 @@
-import inspect
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 from tirex import load_model
-
-from .utils import nanmax, nanmin, nanstd
+from tirex.util import nanmax, nanmin, nanstd
 
 
 class TiRexEmbedding(nn.Module):
-    def __init__(self, device: str | None = None, data_augmentation: bool = False, batch_size: int = 512) -> None:
+    def __init__(
+        self, device: str | None = None, data_augmentation: bool = False, batch_size: int = 512, compile: bool = False
+    ) -> None:
         super().__init__()
         self.data_augmentation = data_augmentation
         self.number_of_patches = 8
         self.batch_size = batch_size
 
         if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
         self.device = device
+        self._compile = compile
 
-        self.model = load_model(path="NX-AI/TiRex", device=self.device)
+        self.model = load_model(path="NX-AI/TiRex", device=self.device, compile=self._compile)
 
     def _gen_emb_batched(self, data: torch.Tensor) -> torch.Tensor:
         batches = list(torch.split(data, self.batch_size))