Table of Contents
Why saving and loading tensors matters
Tensors are the raw material of PyTorch work. You might want to save a preprocessed dataset tensor to avoid repeating expensive preprocessing, cache embeddings, store intermediate results for debugging, or checkpoint a long running computation. In this chapter you will learn the practical mechanics of saving and loading tensors, how to preserve device and dtype correctly, and how to avoid common file format mistakes.
The core APIs: `torch.save` and `torch.load`
PyTorch’s standard way to serialize tensors is torch.save and torch.load. Under the hood this uses Python pickling plus PyTorch’s storage format, which is designed to preserve tensor metadata such as shape, dtype, and device information.
A minimal example looks like this:
import torch
x = torch.randn(3, 4)
torch.save(x, "x.pt")
y = torch.load("x.pt")
print(y.shape, y.dtype, y.device)
torch.save can write a single tensor, or a Python object that contains tensors, such as a dict or list. This is often the most convenient way to bundle multiple related tensors.
Saving multiple tensors together
A common pattern is to save a dictionary of tensors, possibly with a few simple metadata values.
import torch
data = {
"features": torch.randn(100, 20),
"labels": torch.randint(0, 5, (100,)),
"mean": torch.tensor(0.123), # tensors are fine
"std": torch.tensor(0.987),
}
torch.save(data, "dataset_cache.pt")
loaded = torch.load("dataset_cache.pt")
X = loaded["features"]
y = loaded["labels"]You can also save tuples and lists, but dicts are easier to extend later without breaking your loading code.
Only load .pt files you trust. torch.load uses pickle and can execute arbitrary code during loading if the file is malicious.
File extensions and conventions
PyTorch does not require a particular extension, but common conventions are .pt and .pth. You can use either for tensors and PyTorch saved objects. Prefer a consistent naming scheme, for example train_features.pt, embeddings_cache.pt, or checkpoint_step_1000.pt.
Device handling when loading, `map_location`
If you save a CUDA tensor and later load on a machine without a GPU, you must map the tensors to CPU during load. Use map_location.
import torch
# Suppose this was saved from a GPU machine
obj = torch.load("x.pt", map_location="cpu")You can also map to a specific GPU.
obj = torch.load("x.pt", map_location="cuda:0")If you want to load on CPU first and move later, that is also fine and is often a clean approach for portability.
obj = torch.load("x.pt", map_location="cpu")
obj = obj.to("cuda:0")
Do not assume that a saved tensor will load onto the device you want. Use map_location explicitly in code that needs to run on different hardware.
Dtype preservation and intentional casting after load
torch.save preserves dtype. If you need a different dtype for training or inference, cast after loading.
x = torch.load("x.pt", map_location="cpu")
x_fp16 = x.to(torch.float16)
x_fp32 = x.to(torch.float32)This is especially useful when you store tensors in a compact dtype to save disk space, then convert at runtime.
Saving and loading with paths
Use pathlib.Path for robust path handling.
from pathlib import Path
import torch
path = Path("artifacts") / "tensor_cache.pt"
path.parent.mkdir(parents=True, exist_ok=True)
torch.save(torch.arange(10), path)
t = torch.load(path)CPU tensors, GPU tensors, and portability
Saving GPU tensors is allowed, but for maximum portability you can move to CPU before saving.
x = x.detach().cpu()
torch.save(x, "x_cpu.pt")This is a good default when saving data artifacts that will be shared across different environments.
If a tensor requires gradients, saving it will also save enough information to restore it as a tensor with requires_grad set. For cached data and features, you usually want detached tensors. Use x = x.detach() before saving.
Memory considerations when loading large tensors
torch.load will materialize tensors in memory. For very large saved tensors, ensure you have enough RAM or GPU memory for the loaded object. A simple practical tactic is to store separate files for very large components so you can load only what you need, for example features.pt and labels.pt instead of one combined file.
When not to use `torch.save` for tensors
torch.save is excellent for PyTorch workflows, but it is not the best interchange format for non PyTorch environments. If you need to share arrays with other tools, you might prefer NumPy formats. In that case you can convert with tensor.cpu().numpy() and use NumPy saving, but that is outside the scope of this chapter.
Quick checklist
Rule of thumb: save detached CPU tensors for portability, and always use map_location when loading in code that may run on different devices.