Create a Simple Neural Network Class

easyFREEModelingPyTorch ModelingPytorch

Problem Statement

Create a basic neural network class by subclassing nn.Module.

Background

In PyTorch, neural networks are defined by subclassing nn.Module. You define the layers in the constructor and specify how data flows through them in the forward method.

Your Task

Write a function create_simple_network() that defines a single-layer neural network (10 inputs → 5 outputs) and returns a dictionary describing the network.

Output Format

Return a dictionary with keys: "class_name" (string), "is_module" (boolean), "output_shape" (list — shape when given a single sample of 10 features).

Example

Input

None

Output

{'class_name': 'SimpleNet', 'is_module': True, 'output_shape': [1, 5]}

Reasoning

Input: the function takes no arguments. Internally it builds a network and feeds it one test sample: a single row of 10 random features, i.e. a $1 \times 10$ batch (1 sample, 10 features).
The math: a single fully connected layer computes $y = Wx + b$ , where $W$ is a $5 \times 10$ weight matrix and $b$ is a bias of length 5. Each output $y_j = \sum_{i=1}^{10} W_{ji} x_i + b_j$ , so 10 inputs collapse into 5 numbers. That is $5 \times 10 + 5 = 55$ learnable parameters.
Shape: the batch dimension passes through untouched, so $1 \times 10 \to 1 \times 5$ , giving output_shape = [1, 5]. A batch of 32 samples would give $32 \times 5$ .
class_name: the name of the class you defined, SimpleNet.
is_module: True, because the class inherits from PyTorch's module base class, so it is a module.
Notice that the actual output values are random (the input and initial weights are random); only the shape is checked.

Constraints

Must subclass nn.Module
Must call super().init() in init
Use nn.Linear(10, 5)

Learn this in

PyTorch FundamentalsFoundations

More from Pytorch

Apply Softmax for Predictionseasy Count Model Parametersmedium Build with nn.Sequentialeasy Use nn.Flatten Layereasy Compose Multiple Transformsmedium

Concept primer

1. Background Knowledge

PyTorch nn.Module is the foundation for all neural networks. Every neural network in PyTorch inherits from nn.Module, which provides essential functionality like parameter tracking, device placement, and the forward pass mechanism.

Key concepts:

Parameters: nn.Module automatically registers layers like nn.Linear as trainable parameters via self.layer = nn.Linear(...)
Forward pass: Defined in forward() method; called automatically when using model(input)
Linear layer: nn.Linear(in_features, out_features) implements $y = xW^T + b$ where $W \in \mathbb{R}^{out \times in}$ , $b \in \mathbb{R}^{out}$
Shape transformation: Input (batch_size, input_size) → Output (batch_size, output_size)

Prerequisites:

import torch
import torch.nn as nn

2. Algorithm Approach

This is a feedforward neural network with a single linear transformation:

Input: (batch_size, 10) → Linear(10, 5) → Output: (batch_size, 5)

Mathematical operation: \mathbf{y} $=$ \mathbf{x}\mathbf{W} $^T +$ \mathbf{b}

\mathbf{x} $\in \mathbb{R}^{batch \times 10}$
\mathbf{W} $\in \mathbb{R}^{5 \times 10}$ (weights)
\mathbf{b} $\in \mathbb{R}^{5}$ (bias)
\mathbf{y} $\in \mathbb{R}^{batch \times 5}$

3. Step-by-Step Strategy

Define the class inheriting from nn.Module:

class SimpleNet(nn.Module):

Initialize in init():

Call super(SimpleNet, self).init() first
Create self.linear = nn.Linear(10, 5)

Define forward():

Return self.linear(x) (no other operations needed)

Create create_simple_network():

Instantiate SimpleNet()
Verify it's nn.Module subclass: isinstance(net, nn.Module)
Test output shape: net(torch.randn(1, 10)).shape
Return dictionary with all 3 keys

Complete solution structure:

def create_simple_network():
 class SimpleNet(nn.Module): # Defined inside function
 def __init__(self):
 super().__init__()
 self.linear = nn.Linear(10, 5)

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

 net = SimpleNet()
 return {
 "class_name": "SimpleNet",
 "is_module": isinstance(net, nn.Module),
 "output_shape": list(net(torch.randn(1, 10)).shape)
 }

4. Common Pitfalls

Mistake	Why it fails	Fix
Forget super().init()	Parameters not registered	Always call first in init
Call net.forward(x)	Bypasses hooks/autocast	Use net(x)
Wrong shape test	Input must be (1, 10)	torch.randn(1, 10)
Return class instead of dict	Test expects dictionary	Return {"class_name":...,...}
nn.Linear(5, 10)	Wrong dimensions	nn.Linear(10, 5) per spec

Critical: Class must be named exactly "SimpleNet" for the test.

5. Time & Space Complexity

Forward pass: $O(n \cdot d_{in} \cdot d_{out}) = O(1 \cdot 10 \cdot 5) = O(50)$ per sample

Matrix multiplication dominates
Constant time for fixed architecture

Space complexity:

Parameters: $10 \times 5 + 5 = 55$ scalars
Activations: $O(batch\_size \cdot 5)$
Total: $O(1)$ (fixed-size model)

Verification test runs in $O(1)$ time since batch_size=1 and layers are tiny.

This pattern scales to complex networks: more layers → $O($ \sum $L_i \cdot d_i \cdot d_{i+1})$ .

solution.py

Python 3.13

Test results