PyTorch

PyTorch is a tensor and neural-network library. Each PyTorch layer defines the shape, data type, and device that it expects. There is no single axis order for every layer, so inspect the layer interface before supplying a tensor. See TensorFlow for the corresponding Keras conventions.

Tensor Shape Rules

PyTorch normally puts the batch dimension first for feed-forward data:

Data	Common shape
Feature rows	(batch, features)
Images	(batch, channels, height, width)
Token IDs	(batch, time) when the surrounding model uses batch first

Recurrent layers are a special case. nn.RNN, nn.GRU, and nn.LSTM use this default input order:

(time, batch, features)

Set batch_first=True to use:

(batch, time, features)

batch_first does not affect hidden states

It changes the input and sequence-output layout. Hidden and cell states keep (layers * directions, batch, hidden size).

For exact GRU examples, see the Code section in Gated Recurrent Units (GRU).

Layer Arguments

PyTorch layer constructors often require both the input size and output size:

linear = nn.Linear(in_features=8, out_features=4)
gru = nn.GRU(input_size=8, hidden_size=6, batch_first=True)
embedding = nn.Embedding(num_embeddings=1000, embedding_dim=64)

Read these as:

Layer	Input rule	Output rule
nn.Linear(8, 4)	Final input dimension must be 8	Final output dimension becomes 4
nn.GRU(8, 6)	Each time step has 8 input features	Each direction has a hidden size of 6
nn.Embedding(1000, 64)	Integer IDs range from 0 to 999	Each ID becomes 64 values

nn.Linear changes only the final dimension:

(3, 5, 8) -> Linear(8, 4) -> (3, 5, 4)

PyTorch stores Linear.weight with shape (out_features, in_features). The forward operation uses the compatible transpose of that stored weight.

Data Types and Devices

Common rules:

• neural-network activations and model weights normally use floating-point tensors
• embedding inputs and class-index targets normally use torch.long
• a model and the tensors used by that model must be on the same device

x = x.to(device)
model = model.to(device)

Loss Call Order

PyTorch loss objects normally receive predictions first and targets second:

criterion = nn.CrossEntropyLoss()
predictions = model(x)
loss = criterion(predictions, targets)

criterion is a variable name for the loss function. It is not a Python keyword.

Manual Training Update

PyTorch commonly exposes each training step:

optimizer.zero_grad()
predictions = model(x)
loss = criterion(predictions, targets)
loss.backward()
torch.nn.utils.clip_grad_value_(model.parameters(), 0.5)
optimizer.step()

The order is:

clear old gradients
-> run the forward pass
-> calculate the loss
-> calculate gradients with backpropagation
-> clip the calculated gradients
-> update the parameters

backward() calculates and stores gradients in parameter.grad. step() reads those gradients and changes the parameters.

Joining and Rearranging Tensors

torch.cat joins tensors along an existing dimension:

a = torch.zeros(3, 5)
b = torch.zeros(3, 5)
 
torch.cat((a, b), dim=1).shape  # (3, 10)
torch.cat((a, b), dim=0).shape  # (6, 5)

Related operations have different purposes:

Operation	Purpose
torch.cat	Join along an existing dimension
torch.stack	Join while creating a new dimension
reshape	Change the dimension sizes without changing value order
permute	Change the axis order

The official interfaces are listed in the PyTorch documentation.