Skip to main content

Overview

The torch package contains data structures for multi-dimensional tensors and defines mathematical operations over these tensors. It provides utilities for efficient serialization and CUDA support for GPU computation.

Core Functions

Tensor Creation

Creates a tensor from data.
data
array_like
required
Initial data for the tensor. Can be a list, tuple, NumPy ndarray, scalar, and other types.
dtype
torch.dtype
The desired data type of returned tensor. If None, infers data type from data.
device
torch.device
The desired device of returned tensor. Default: CPU for CPU tensor types, current CUDA device for CUDA tensor types.
requires_grad
bool
default:"False"
If autograd should record operations on the returned tensor.
pin_memory
bool
default:"False"
If set, returned tensor would be allocated in pinned memory. Works only for CPU tensors.
import torch

# Create tensor from list
x = torch.tensor([[1, 2], [3, 4]])

# Create with specific dtype
y = torch.tensor([1.0, 2.0], dtype=torch.float32)

# Create on CUDA device
z = torch.tensor([1, 2, 3], device='cuda')
Returns a tensor filled with the scalar value 0.
*size
int
required
A sequence of integers defining the shape of the output tensor.
dtype
torch.dtype
The desired data type. Default: global default dtype.
device
torch.device
The desired device of returned tensor.
requires_grad
bool
default:"False"
If autograd should record operations on the returned tensor.
# Create 3x4 zero matrix
x = torch.zeros(3, 4)
# tensor([[0., 0., 0., 0.],
#         [0., 0., 0., 0.],
#         [0., 0., 0., 0.]])
Returns a tensor filled with the scalar value 1.
*size
int
required
A sequence of integers defining the shape of the output tensor.
dtype
torch.dtype
The desired data type. Default: global default dtype.
device
torch.device
The desired device of returned tensor.
x = torch.ones(2, 3)
# tensor([[1., 1., 1.],
#         [1., 1., 1.]])
Returns a tensor filled with random numbers from a uniform distribution on the interval [0, 1).
*size
int
required
A sequence of integers defining the shape of the output tensor.
generator
torch.Generator
A pseudorandom number generator for sampling.
dtype
torch.dtype
The desired data type. Default: global default dtype.
device
torch.device
The desired device of returned tensor.
x = torch.rand(2, 3)
# tensor([[0.1234, 0.5678, 0.9012],
#         [0.3456, 0.7890, 0.2345]])
Returns a tensor filled with random numbers from a normal distribution with mean 0 and variance 1.
*size
int
required
A sequence of integers defining the shape of the output tensor.
generator
torch.Generator
A pseudorandom number generator for sampling.
dtype
torch.dtype
The desired data type. Default: global default dtype.
x = torch.randn(2, 3)
# tensor([[-0.5423,  0.2347, -1.3478],
#         [ 1.5678, -0.7890,  0.4567]])

Mathematical Operations

Computes the absolute value of each element in input.
input
Tensor
required
The input tensor.
out
Tensor
The output tensor.
Formula: out_i = |input_i|
>>> torch.abs(torch.tensor([-1, -2, 3]))
tensor([1, 2, 3])
Adds other, scaled by alpha, to input.
input
Tensor
required
The first input tensor.
other
Tensor
required
The second input tensor or scalar.
alpha
Number
default:"1"
The multiplier for other.
Formula: out = input + alpha × other
>>> a = torch.tensor([1, 2, 3])
>>> b = torch.tensor([4, 5, 6])
>>> torch.add(a, b, alpha=2)
tensor([9, 12, 15])
Matrix product of two tensors with broadcasting support.
input
Tensor
required
The first tensor to be multiplied.
other
Tensor
required
The second tensor to be multiplied.
out
Tensor
The output tensor.
The behavior depends on the dimensionality of the tensors:
  • If both tensors are 1-dimensional, the dot product (scalar) is returned.
  • If both arguments are 2-dimensional, the matrix-matrix product is returned.
  • If the first argument is 1-dimensional and the second argument is 2-dimensional, a 1 is prepended to its dimension for the purpose of the matrix multiply. After the matrix multiply, the prepended dimension is removed.
  • If the first argument is 2-dimensional and the second argument is 1-dimensional, the matrix-vector product is returned.
# Matrix multiplication
>>> a = torch.randn(2, 3)
>>> b = torch.randn(3, 4)
>>> torch.matmul(a, b).shape
torch.Size([2, 4])

# Batch matrix multiplication
>>> a = torch.randn(10, 3, 4)
>>> b = torch.randn(10, 4, 5)
>>> torch.matmul(a, b).shape
torch.Size([10, 3, 5])
Returns the sum of all elements in the input tensor.
input
Tensor
required
The input tensor.
dim
int or tuple of ints
The dimension or dimensions to reduce. If None, all dimensions are reduced.
keepdim
bool
default:"False"
Whether the output tensor has dim retained or not.
dtype
torch.dtype
The desired data type of returned tensor. If specified, the input tensor is casted to dtype before operation.
>>> a = torch.randn(4, 4)
>>> torch.sum(a)
tensor(-2.4567)

>>> torch.sum(a, dim=0)
tensor([-0.5, 1.2, -0.8, 0.3])

Utilities

Saves an object to a disk file.
obj
object
required
The saved object (can be a tensor, model state_dict, or any picklable Python object).
f
str or file-like object
required
A file-like object (has to implement write and flush) or a string containing a file name.
# Save a tensor
x = torch.tensor([1, 2, 3])
torch.save(x, 'tensor.pt')

# Save a model's state dict
torch.save(model.state_dict(), 'model_weights.pth')
Loads an object saved with torch.save() from a file.
f
str or file-like object
required
A file-like object or a string containing a file name.
map_location
function, torch.device, string or dict
A function, torch.device, string or dict specifying how to remap storage locations.
weights_only
bool
default:"False"
If True, only loads tensors, primitive types, and dictionaries.
# Load a tensor
x = torch.load('tensor.pt')

# Load model weights
model.load_state_dict(torch.load('model_weights.pth'))

# Load to specific device
model = torch.load('model.pth', map_location='cuda:0')
Sets the seed for generating random numbers. Returns a torch.Generator object.
seed
int
required
The desired seed value.
torch.manual_seed(42)
x = torch.randn(2, 3)
# Results are now reproducible
Sets the default torch.Tensor device.
device
torch.device or str
required
The device to set as default.
torch.set_default_device('cuda')
x = torch.randn(2, 3)  # Now on CUDA by default
Sets the internal precision of float32 matrix multiplications.
precision
str
required
Can be “highest”, “high”, or “medium”.
Running float32 matrix multiplications in lower precision may significantly increase performance:
  • “highest”: float32 datatype (24 mantissa bits) for internal computations
  • “high”: TensorFloat32 datatype (10 mantissa bits) or bfloat16-based algorithm
  • “medium”: bfloat16 datatype (8 mantissa bits) for internal computations
torch.set_float32_matmul_precision('high')
# Float32 matmuls now use TensorFloat32 for speed

Configuration

Sets whether PyTorch operations must use deterministic algorithms.
mode
bool
required
If True, operations switch to deterministic algorithm or throw runtime error.
warn_only
bool
default:"False"
If True, operations without deterministic implementation throw warning instead of error.
torch.use_deterministic_algorithms(True)
# Now all operations use deterministic algorithms
Returns True if obj is a PyTorch tensor.
obj
object
required
Object to test.
return
bool
True if obj is a PyTorch tensor.
>>> torch.is_tensor(torch.tensor([1, 2, 3]))
True
>>> torch.is_tensor([1, 2, 3])
False

Symbolic Types

SymInt, SymFloat, SymBool

PyTorch provides symbolic types for dynamic shape tracking and symbolic computation:
import torch

# SymInt tracks symbolic integer values
x = torch.randn(torch.SymInt(10), 20)
print(x.shape)  # torch.Size([s0, 20])

Best Practices

  • Use torch.empty() when you’ll immediately fill the tensor
  • Use torch.zeros_like() and torch.ones_like() to match existing tensor properties
  • Specify device='cuda' directly during creation to avoid CPU→GPU copy
Set all random seeds for reproducible results:
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)
  • Use torch.no_grad() context when not training to reduce memory usage
  • Call del on large tensors and torch.cuda.empty_cache() to free GPU memory
  • Use gradient checkpointing for large models

Tensor API

Tensor class methods and operations

Autograd API

Automatic differentiation functions

CUDA API

GPU acceleration with CUDA

Build docs developers (and LLMs) love