Code360 powered by Coding Ninjas X Code360 powered by Coding Ninjas X
Last Updated: Mar 27, 2024
Difficulty: Easy

torch.nn Module in PyTorch

Leveraging ChatGPT - GenAI as a Microsoft Data Expert
Prerita Agarwal
Data Specialist @
23 Jul, 2024 @ 01:30 PM


A neural network consists of algorithms that aim to understand and recognize the hidden patterns in raw data, which can be used to extract meaningful information and solve problems in artificial intelligence like medical diagnosis, credit card fraud detection, etc. The PyTorch framework helps develop the neural network, train, and build the model quickly. It has high-level APIs to build neural networks.

What is the torch.nn module?

PyTorch contains torch.nn module is used to train and build the layers of neural networks such as input, hidden, and output. Torch.nn base class helps wrap the torch's parameters, functions, and layers.nn module.

Get the tech career you deserve, faster!
Connect with our expert counsellors to understand how to hack your way to success
User rating 4.7/5
1:1 doubt support
95% placement record
Akash Pal
Senior Software Engineer
326% Hike After Job Bootcamp
Himanshu Gusain
Programmer Analyst
32 LPA After Job Bootcamp
After Job

Classes in the torch.nn module

Torch.nn contains various classes and modules. Some of them are:

  • Parameters
    This torch.nn.Parameter() subclass can store learnable initial states and hidden states of models.
  • Containers
    This container class uses nn.Container() subclass to create complex neural networks.
    Examples - torch.nn.Sequential() is used to combine different layers, torch.nn.ParameterDict() to store the parameters in a dictionary, etc.
  • Layers
    We can configure various trainable layers in a neural network using torch.nn.
    Examples - Padding layers add value to the sides of a tensor, Recurrent layers, Sparse layers, etc. 
  • Functions
    Torch.nn module contains various loss functions to evaluate the error between the input and the target values. Example - torch.nn.L1Loss() is used to calculate the mean absolute error between input and output, torch.nn.CrossEntropyLoss(), etc.
    Torch.nn also contains functions to calculate the distance between two parameters, such as torch.nn.CosineSimilarity() to calculate the cosine similarity between two variables.

Please refer to this link to learn more about these classes and modules.

Implementation of torch.nn functions

Let's build some neural network architectures using the torch.nn module.

import torch
import torch.nn.functional as F

#Sample network

class Net(torch.nn.Module):

    def __init__(self, n_feature, n_hidden, n_output):

        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_feature, n_hidden) #hidden layer
        self.predict = torch.nn.Linear(n_hidden, n_output) #output layer

    def forward(self, x):

        x = F.relu(self.hidden(x))  #relu activation function for hidden layer
        x = self.predict(x)         #output
        return x

network_1 = Net(2, 11, 2) #example



  (hidden): Linear(in_features=2, out_features=11, bias=True)
  (predict): Linear(in_features=11, out_features=2, bias=True)


Another example:

#faster way to build neural network

network_2 = torch.nn.Sequential(
    torch.nn.Linear(2, 11),
    torch.nn.Linear(11, 2)




  (0): Linear(in_features=2, out_features=11, bias=True)
  (1): ReLU()
  (2): Linear(in_features=11, out_features=2, bias=True)


Let us see how to use the cross-entropy loss function using the torch.nn module.

from torch import nn

loss_fun = nn.CrossEntropyLoss()
input = torch.tensor([[0.1,0.2,0.3,0.4]],dtype=torch.float) #sample input
target = torch.tensor([0], dtype=torch.long) #sample output
loss_fun(input, target)





Now calculate the cosine similarity between two random tensors using the torch.nn module.

import torch
import torch.nn.functional as F

tensor1 = torch.randn(50) #random tensor of size=50
tensor2 = torch.randn(50) #random tensor of size=50

cosine_similarity_value = F.cosine_similarity(tensor1, tensor2, dim=0, eps=1e-6) 
#eps is a small value to avoid division by zero 															 





Frequently Asked Questions

1. What is the torch.nn module in Python?
The torch.nn module helps in developing and building neural networks quickly.

2. What is nn.linear in PyTorch?
nn.linear(n,m) module takes n inputs to create a single-layer feed-forward network with m outputs.

3. What is PyTorch?
PyTorch is a library in Python that helps in developing neural networks conveniently.

4. How can you calculate cross-entropy loss using the torch.nn module?
We can use the torch.nn.CrossEntropyLoss() function of torch.nn for calculation.

5. What does nn.sequential do?
It helps to run layers sequentially quickly.


This article discussed the torch.nn module in PyTorch, its uses, and the implementation of various functions and methods present in it.

Check out this article - Padding In Convolutional Neural Network

We hope this blog has helped you enhance your knowledge regarding torch.nn module in PyTorch, and if you would like to learn more, check out our free content on Machine Learning and more unique courses. Do upvote our blog to help other ninjas grow.

Happy Coding!

Topics covered
What is the torch.nn module?
Classes in the torch.nn module
Implementation of torch.nn functions
Frequently Asked Questions