Skip to content
PyTorch

Tensor Basics

Create, index, and operate on tensors.

By EZ4Code Team
tensorbasics

Code

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)

Explanation

Tensors are multi-dimensional arrays similar to NumPy ndarrays but with GPU support and automatic differentiation. Constructors like zeros, ones, and randn mirror NumPy, and operations like sum and matmul are available as methods or functions. to() moves tensors between CPU and GPU devices.

More PyTorch Snippets