🎓 A1Learner
🔍

PyTorch Type Casting Tensors

Type casting changes a tensor's dtype, the kind of number each cell holds. Methods like .float() and .long() return a new tensor in the new dtype.

Cast with .float() and .long()

Each cast returns a new tensor. The prints show the dtype before and after.

counts = torch.tensor([3, 7, 12])
scores = counts.float()
ids = torch.tensor([1.9, 2.1]).long()
print("Counts:", counts.dtype)
print("As float:", scores.dtype, scores.tolist())
print("As long:", ids.dtype, ids.tolist())

A float model refuses long inputs

A Linear layer holds float weights. Feeding it a long tensor fails, and the last line of the error names both dtypes.

model = nn.Linear(3, 2)
counts = torch.tensor([[1, 2, 3], [4, 5, 6]])
model(counts)

Cast with .to(dtype)

The .to() method takes the target dtype as an argument. The original tensor keeps its dtype.

target_dtype = torch.float32
counts = torch.tensor([3, 7, 12])
cast = counts.to(target_dtype)
print("Cast dtype:", cast.dtype)
print("Original dtype:", counts.dtype)

Example

import torch
from torch import nn

features = torch.tensor([[2, 0], [0, 3], [3, 1], [1, 2]])
labels = torch.tensor([0.0, 1.0, 0.0, 1.0])

x = features.float()
y = labels.long()

torch.manual_seed(0)
model = nn.Linear(2, 2)
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(model(x), y)

print("Features dtype:", x.dtype)
print("Labels dtype:", y.dtype)
print("Labels:", y.tolist())
print("Loss:", round(loss.item(), 4))

Key points

  • Use .float() to cast a tensor to float32
  • Use .long() to cast a tensor to int64
  • Use .to(dtype) to cast to a dtype named in code
  • Casting returns a new tensor and leaves the original unchanged
  • A layer and its input must share a dtype
  • Each loss names the dtype its targets must have

In this example, the features start as whole-number counts. Their dtype is int64. The labels start as float class numbers. The model needs float features, so .float() casts them. Cross-entropy needs long targets, so .long() casts the labels. The .long() cast drops the decimal part of each label. Seed 0 is set before the model is built. The cast features run through the Linear layer. The loss accepts the cast labels and prints 1.2262. Two casts gave the model and the loss the dtypes each one demands.

Note

.float() is a short name for a cast to float32

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A dataset stores class labels as float numbers. Cross-entropy needs them as int64. Which call makes the right targets?