# How to Load a Locally Saved TimesFM Model: PyTorch and JAX Guide

> Learn to load a locally saved TimesFM model using PyTorch and JAX. Follow our guide to instantiate TimesFmCheckpoint and TimesFm for seamless model integration.

- Repository: [Google Research/timesfm](https://github.com/google-research/timesfm)
- Tags: how-to-guide
- Published: 2026-04-02

---

**To load a locally saved TimesFM model, instantiate a `TimesFmCheckpoint` dataclass with the `path` attribute set to your local checkpoint file, ensure your `TimesFmHparams` match the saved configuration, and pass both to the `TimesFm` constructor.**

The google-research/timesfm library provides a unified forecast API for both JAX and PyTorch backends. When you need to load a locally saved TimesFM model—whether downloaded previously, fine-tuned, or transferred between environments—the library offers a straightforward checkpoint loading mechanism that bypasses Hugging Face Hub entirely.

## Understanding the TimesFM Loading Architecture

### The Checkpoint Dataclass

In [`v1/src/timesfm/timesfm_base.py`](https://github.com/google-research/timesfm/blob/main/v1/src/timesfm/timesfm_base.py) (lines 184-203), the library defines `TimesFmCheckpoint`, a dataclass that holds checkpoint location and metadata. When the `path` attribute is provided, the library reads directly from disk without making network calls to Hugging Face.

### Backend-Specific Implementations

The `TimesFm` class acts as an entry point that resolves to `TimesFmJax` or `TimesFmTorch` depending on your runtime. According to the source code:

- **PyTorch**: `TimesFmTorch.load_from_checkpoint` in [`v1/src/timesfm/timesfm_torch.py`](https://github.com/google-research/timesfm/blob/main/v1/src/timesfm/timesfm_torch.py) (lines 52-70) uses `torch.load` or `safetensors.load_file` when a local path is provided.
- **JAX**: `TimesFmJax.load_from_checkpoint` in [`v1/src/timesfm/timesfm_jax.py`](https://github.com/google-research/timesfm/blob/main/v1/src/timesfm/timesfm_jax.py) (lines 94-107) handles Flax checkpoints using Pax utilities.

## Loading a Locally Saved TimesFM Model in PyTorch

Most users interact with the PyTorch backend. The concrete architecture for the 2.5 model family lives in [`src/timesfm/timesfm_2p5/timesfm_2p5_torch.py`](https://github.com/google-research/timesfm/blob/main/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py).

### Standard PyTorch Checkpoint (.ckpt)

```python
from timesfm import TimesFm, TimesFmHparams, TimesFmCheckpoint

hparams = TimesFmHparams(
    per_core_batch_size=32,
    horizon_len=128,
    context_len=512,
    num_layers=20,
    num_heads=16,
    model_dims=1280,
    backend="gpu",
    use_positional_embedding=True,
)

ckpt = TimesFmCheckpoint(
    version="torch",
    path="/absolute/path/to/torch_model.ckpt",
)

model = TimesFm(hparams=hparams, checkpoint=ckpt)

```

The constructor automatically invokes `load_from_checkpoint`, which detects the provided path and loads via `torch.load` instead of calling `snapshot_download`.

### Safetensors Checkpoint

For checkpoints saved in the Hugging Face Safetensors format (recommended for size efficiency), you can load manually as demonstrated in [`v1/src/finetuning/finetuning_example.py`](https://github.com/google-research/timesfm/blob/main/v1/src/finetuning/finetuning_example.py) (lines 152-168):

```python
from timesfm import TimesFm, TimesFmHparams, TimesFmCheckpoint
from safetensors.torch import load_file

hparams = TimesFmHparams(
    per_core_batch_size=32,
    horizon_len=128,
    # Match your checkpoint configuration

)

tmp = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint())
state_dict = load_file("/absolute/path/to/model.safetensors")
tmp._model.load_state_dict(state_dict)
tmp._model.eval()
model = tmp

```

## Loading a Locally Saved TimesFM Model in JAX

For JAX/Flax backends, the process is similar but targets directories containing Flax checkpoints:

```python
from timesfm import TimesFm, TimesFmHparams, TimesFmCheckpoint

hparams = TimesFmHparams(
    per_core_batch_size=8,
    horizon_len=256,
    backend="gpu",
)

ckpt = TimesFmCheckpoint(
    version="jax",
    path="/absolute/path/to/checkpoints",  # Directory containing FLAX checkpoints

)

model = TimesFm(hparams=hparams, checkpoint=ckpt)

```

## Summary

- **Create matching hyperparameters**: Your `TimesFmHparams` must exactly match the architecture of the saved checkpoint (layers, dimensions, horizon length).
- **Use `TimesFmCheckpoint` with `path`**: Setting the `path` attribute to a local file or directory prevents automatic Hugging Face downloads according to the source code in [`timesfm_torch.py`](https://github.com/google-research/timesfm/blob/main/timesfm_torch.py).
- **Backend-specific paths**: PyTorch accepts `.ckpt` or `.safetensors` files, while JAX expects directories containing Flax checkpoint files.
- **Manual loading for Safetensors**: When working with `.safetensors` files, use `load_file` and `load_state_dict` as shown in the fine-tuning example.

## Frequently Asked Questions

### Does TimesFM download from Hugging Face if I provide a local path?

No. When the `path` attribute of `TimesFmCheckpoint` is provided, the `load_from_checkpoint` method in both [`timesfm_torch.py`](https://github.com/google-research/timesfm/blob/main/timesfm_torch.py) and [`timesfm_jax.py`](https://github.com/google-research/timesfm/blob/main/timesfm_jax.py) reads directly from disk using `torch.load` or Flax utilities without calling `snapshot_download`.

### What file formats does TimesFM support for local loading?

The PyTorch backend supports standard PyTorch checkpoints (`.ckpt`) and Hugging Face Safetensors (`.safetensors`). The JAX backend expects directories containing Flax checkpoint files. As implemented in [`timesfm_torch.py`](https://github.com/google-research/timesfm/blob/main/timesfm_torch.py) (lines 52-70), the loader automatically detects and handles these formats.

### How do I ensure my hyperparameters match the checkpoint?

You must manually configure `TimesFmHparams` with the same values used during training: `num_layers`, `model_dims`, `num_heads`, `horizon_len`, and `context_len`. Mismatched parameters will cause shape errors when loading state dictionaries.

### Can I load a fine-tuned TimesFM model locally?

Yes. After fine-tuning and saving locally, create a `TimesFmCheckpoint` pointing to your saved weights. The fine-tuning example in [`v1/src/finetuning/finetuning_example.py`](https://github.com/google-research/timesfm/blob/main/v1/src/finetuning/finetuning_example.py) demonstrates this pattern by loading Safetensors files and attaching them via `load_state_dict`.