torch¶

def load(self) -> None:
    """Load TorchScript model and prepare for inference.

    Performs:
    - Load model from disk
    - Setup device
    - Move model to device
    - Set evaluation mode
    - Apply precision
    - Warmup inference

    Raises:
        FileNotFoundError: If model file does not exist.
        RuntimeError: If device is not available.
    """
    logger.info(f"Loading model from {self.model_path}")

    # Load model
    self.model = torch.jit.load(str(self.model_path))

    # Setup device (reuse mixin)
    self._torch_device = self._setup_torch_device()

    # Configure model
    self.model.to(self._torch_device)
    self.model.eval()
    self.model = self._apply_precision_to_model(self.model)

    logger.info(f"Model loaded on {self._torch_device}")

    # Warmup
    self._warmup()

def infer(self, input: torch.Tensor) -> R:
    """Run inference on a single input.

    Automatically handles:
    - Moving input to correct device
    - Converting to correct precision
    - Adding batch dimension (if configured)
    - Removing batch dimension (if added)

    Args:
        input: Input tensor. Can be 3D (C, H, W) if auto_add_batch_dim=True,
            or 4D (1, C, H, W) otherwise.

    Returns:
        Model output.  Type depends on model architecture (tensor or tuple).

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        with runtime:
            # 3D input (auto_add_batch_dim=True)
            input = torch.randn(3, 224, 224)
            output = runtime.infer(input)

            # 4D input
            input = torch.randn(1, 3, 224, 224)
            output = runtime.infer(input)
        ```
    """
    if self.model is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Prepare input (reuse mixin)
    input = self._prepare_input(input, self._torch_device)
    input, added_batch = self._add_batch_dim_if_needed(input, self.auto_add_batch_dim)

    # Inference
    with torch.no_grad():
        output = self.model(input)

    # Post-process (reuse mixin)
    return self._remove_batch_dim_if_added(output, added_batch)

def infer_batch(self, inputs: list[torch.Tensor]) -> list[R]:
    """Run inference on a batch of inputs.

    Concatenates inputs into a single batch tensor for efficient
    processing, then splits the output back into individual results.

    Args:
        inputs: List of input tensors. Each should have shape (1, C, H, W).

    Returns:
        List of outputs, one per input. Each maintains batch dimension.

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        with runtime:
            batch = [
                torch.randn(1, 3, 224, 224),
                torch.randn(1, 3, 224, 224),
                torch.randn(1, 3, 224, 224),
            ]

            # Efficient batch processing
            outputs = runtime.infer_batch(batch)

            # outputs[0], outputs[1], outputs[2]
        ```
    """
    if self.model is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Concatenate
    batch = torch.cat(inputs, dim=0).to(self._torch_device)

    if self.precision == Precision.FP16:
        batch = batch.half()

    # Inference
    with torch.no_grad():
        batch_output = self.model(batch)

    # Split (reuse mixin)
    return self._split_batch_output(batch_output, len(inputs))

def unload(self) -> None:
    """Unload model and free resources.

    Performs:
    - Release model from memory
    - Clear CUDA cache (if using CUDA)

    Safe to call multiple times.
    """
    logger.info("Unloading model")
    self.model = None

    if self.device.type.value == "cuda":
        torch.cuda.empty_cache()

    logger.info("Model unloaded")