Skip to content
PyTorch

Tensor Basics

Create, index, and operate on tensors.

#tensor#basics

Code

pytorch
import torch

# Creation
a = torch.tensor([1, 2, 3], dtype=torch.float32)
b = torch.zeros(2, 3)
c = torch.ones(3, 3)
d = torch.randn(2, 2)
e = torch.arange(0, 10, 2).reshape(2, -1)

# Operations
print(a + a, a * 2, a @ a)
print(a.sum(), a.mean(), a.max())

# Reshape, view, and move between devices
flat = d.view(-1)
moved = a.to("cuda" if torch.cuda.is_available() else "cpu")

# Indexing
print(e[:, 0], e[0, :])

# Conversion to and from NumPy
import numpy as np
np_arr = a.numpy()
back = torch.from_numpy(np_arr)