Kahibaro
Discord Login Register

2.2 Tensor Shapes, Dtypes, and Devices

Why shapes, dtypes, and devices matter

In PyTorch, every tensor has a shape, a dtype, and a device. Most beginner bugs come from mixing these up, such as multiplying tensors with incompatible shapes, feeding a model inputs with the wrong dtype, or accidentally keeping some tensors on CPU while the model is on GPU. Getting comfortable reading and controlling these three properties makes everything else in deep learning much smoother.

Tensor shapes

A tensor’s shape tells you how many dimensions it has and how large each dimension is. You can inspect it with x.shape or x.size(). A single number like torch.tensor(3.0) has shape torch.Size([]), which is a scalar. A vector like torch.randn(5) has shape torch.Size([5]). A matrix like torch.randn(2, 3) has shape torch.Size([2, 3]). For images, you commonly see 4D tensors, often shaped like (batch, channels, height, width) in PyTorch.

Many operations require specific shape relationships. Matrix multiplication, for example, follows linear algebra rules: if A is shape (m, n) and B is shape (n, p), then A @ B is shape (m, p). For elementwise operations like A + B, the shapes must match or be broadcastable, which is covered in the broadcasting chapter.

You will often reshape tensors. The most common tools are reshape, view, unsqueeze, squeeze, and permute. reshape changes the shape while keeping the same number of elements, and it can return a copy if needed. view is similar but requires the tensor to be contiguous in memory. unsqueeze adds a dimension of size 1, and squeeze removes dimensions of size 1. permute reorders dimensions, which is common when switching between different layout conventions.

A reshape must preserve the number of elements. If x.numel() changes, the reshape is invalid.
view can fail on non contiguous tensors, such as after permute. In that case, use reshape or call x.contiguous() before view.

A quick way to sanity check shape logic is to print intermediate shapes during development. When a runtime error happens, PyTorch error messages usually include the conflicting shapes. Treat those numbers as clues about which dimension you accidentally swapped or forgot to include, such as a missing batch dimension.

Dtypes in PyTorch

A tensor’s dtype controls how its values are stored and computed. You can inspect it with x.dtype. Common dtypes include torch.float32, torch.float64, torch.float16, torch.bfloat16, and integer types like torch.int64.

Deep learning models almost always use floating point types for inputs and weights, typically float32. Many PyTorch layers and losses expect floating point inputs. Class labels for classification are commonly stored as integer types, often int64, because loss functions like cross entropy treat labels as class indices.

You can set dtype when creating tensors, such as torch.tensor([1, 2], dtype=torch.float32), or convert later with x.to(torch.float32) or x.float(). Conversions create a new tensor, so if you convert frequently in a loop, it can slow things down.

Do not confuse integer and floating point tensors. If you perform operations that require gradients, your tensors must be floating point, because integer tensors cannot require gradients in PyTorch.
Many classification losses expect labels to be integer class indices, typically torch.int64, not one hot float vectors, unless the loss explicitly says so.

Numerical precision also matters. float64 is more precise but slower and uses more memory, and is not usually necessary for neural networks. Mixed precision uses float16 or bfloat16 for speed, mostly on GPUs, but it needs care and is introduced later in the training practice section.

Devices in PyTorch

A tensor’s device tells you where it lives, such as CPU or GPU. You can inspect it with x.device. A CPU tensor has a device like cpu. A CUDA GPU tensor has a device like cuda:0. On Apple Silicon, you may see mps:0 when using the MPS backend.

You move tensors and models to a device with .to(device). A common pattern is to create a device once, then move everything consistently.

When you do an operation, all participating tensors must be on the same device. If one tensor is on CPU and another is on GPU, PyTorch will raise an error instead of silently copying, which is good because silent copies would hide performance problems.

All tensors involved in an operation must be on the same device. This includes model parameters, inputs, targets, and any tensors you create inside the training loop.
A common bug is creating a new tensor on CPU by default inside a loop, such as torch.zeros(...), while your model is on GPU. Prefer creating tensors on the same device, for example by using device=device or by using existing tensors as a reference, such as torch.zeros_like(x).

Another subtlety is that .to(device) returns a tensor on the new device, it does not modify the original tensor in place. The same is true for models, where model.to(device) modifies the module parameters in place, but when moving ordinary tensors you typically write x = x.to(device).

Putting them together in a consistent workflow

In practice, you want your data tensors to have the right shape for the layer you are calling, use floating point dtype for model inputs, and live on the same device as the model. Labels should be the dtype expected by your loss. The simplest way to stay consistent is to choose a device at the start, move the model to that device once, and in each iteration move the batch tensors to that device before the forward pass. If you need to create new tensors during computation, create them on the correct device and with a compatible dtype.

Quick inspection and debugging checklist

When something fails, check x.shape, x.dtype, and x.device for the inputs, the targets, and one or two model parameters. If shapes look wrong, look for missing batch dimensions, swapped axes, or incorrect reshapes. If dtypes look wrong, confirm that your inputs are floating point and your labels match the loss expectations. If devices mismatch, ensure everything has been moved to the same device and that tensors created inside the loop are also created there.

Views: 65

Comments

Please login to add a comment.

Don't have an account? Register now!