Pytorch Python Library CheatSheet

PyTorch is a popular open-source machine learning library known for its flexibility and dynamic computational graph. Whether you’re a seasoned deep learning practitioner or just getting started, having a PyTorch cheatsheet at your fingertips can be incredibly handy. In this blog post, we’ll provide a concise reference guide covering some essential PyTorch concepts and commands.


Tensors are the fundamental building blocks in PyTorch, analogous to arrays in NumPy. They can be scalars, vectors, matrices, or higher-dimensional arrays.

Creating Tensors

import torch

# Create a tensor with random values
tensor_random = torch.rand(3, 4)

# Create a tensor of zeros
tensor_zeros = torch.zeros(2, 2)

# Create a tensor from a list
tensor_list = torch.tensor([1, 2, 3, 4])

Operations on Tensors

# Element-wise addition
result = tensor_random + tensor_zeros

# Matrix multiplication
result_matrix_mul =, tensor_zeros.T)

# Reshaping a tensor
reshaped_tensor = tensor_list.view(2, 2)


PyTorch’s automatic differentiation engine, known as Autograd, enables gradient-based optimization.

# Tensor with gradient tracking
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# Perform operations
y = x.sum()

# Compute gradients

# Access gradients
gradient = x.grad

Neural Networks

Defining and training neural networks in PyTorch involves creating a custom class that inherits from torch.nn.Module.

import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)


Optimizers are used to update the weights of a neural network during training.

import torch.optim as optim

# Define a model and an optimizer
model = SimpleNN()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Inside the training loop
optimizer.zero_grad()  # Clear previous gradients
output = model(input_data)
loss = loss_function(output, target)
loss.backward()  # Compute gradients
optimizer.step()  # Update weights

Loading and Saving Models

Saving and loading trained models is crucial for deployment and sharing.

# Save a model, 'model.pth')

# Load a model
loaded_model = SimpleNN()

GPU Acceleration

Utilizing a GPU can significantly speed up deep learning computations.

# Check if GPU is available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Move tensors and models to GPU
tensor_on_gpu =
model_on_gpu = SimpleNN().to(device)

This PyTorch cheatsheet provides a quick reference for common operations and concepts. While this is by no means an exhaustive guide, it should serve as a handy reference for both beginners and experienced practitioners. As you delve deeper into PyTorch, exploring the official documentation and additional resources is recommended to unlock the full potential of this powerful deep learning library.


1. What is the key difference between PyTorch and TensorFlow?

PyTorch and TensorFlow are both popular deep learning frameworks, but they differ in their approach to building and training models. PyTorch is known for its dynamic computational graph, allowing for more flexibility and easier debugging during model development. TensorFlow, on the other hand, uses a static computational graph, which can lead to more efficient deployment in certain scenarios. The choice between them often depends on personal preference and the specific requirements of the project.

2. How can I check if my machine has a compatible GPU for PyTorch?

PyTorch supports GPU acceleration, which can significantly speed up training times for deep learning models. To check if your machine has a compatible GPU, you can use the following PyTorch code:
import torch # Check if GPU is available if torch.cuda.is_available(): print("GPU is available!") else: print("No GPU detected; using CPU.")

3. What is Autograd in PyTorch, and why is it important?

Autograd (automatic differentiation) is a crucial feature in PyTorch that automatically computes gradients of tensors during backpropagation. This enables the implementation of gradient-based optimization algorithms for training neural networks. By tracking operations on tensors and calculating their gradients, Autograd simplifies the process of updating model parameters and is a key component of PyTorch’s ease of use.

4. How can I load a pre-trained model in PyTorch?

Loading a pre-trained model in PyTorch involves creating an instance of your custom model architecture and loading the pre-trained weights. Here’s a simple example:
import torch import torch.nn as nn class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() # ... define your model architecture ... # Instantiate your model model = MyModel() # Load pre-trained weights model.load_state_dict(torch.load('pretrained_model.pth'))

5. What is the purpose of zero_grad() in PyTorch optimizers?

The zero_grad() method in PyTorch optimizers is used to clear the gradients of all model parameters. During the training loop, gradients accumulate as you perform backpropagation. Calling zero_grad() before each iteration ensures that the gradients from the previous iteration do not interfere with the current one. This is a critical step in gradient-based optimization to prevent unintentional gradient accumulation and incorrect parameter updates.