Bug Report for handwritten-digit-classifier

[saptarshi059]

Bug Report for https://neetcode.io/problems/handwritten-digit-classifier

Please describe the bug below and include any steps to reproduce the bug or screenshots if possible.

This is my code for this problem:

import torch
import torch.nn as nn
from torchtyping import TensorType

class Solution(nn.Module):
    def __init__(self):
        super().__init__()
        torch.manual_seed(0)
        # Define the architecture here
        
        # first use linear layer with 512 neurons. This will map images of size 784 to 512.
        self.ff1 = nn.Linear(784, 512)

        # ReLU activation
        self.relu = nn.ReLU()

        # dropout layer with probability p = 0.2
        self.dropout = nn.Dropout(p = 0.2)

        # final Linear layer with 10 neurons
        self.ff2 = nn.Linear(512, 10)

        # sigmoid activation
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, images: TensorType[float]) -> TensorType[float]:
        torch.manual_seed(0)
        # Return the model's prediction to 4 decimal places
        op1 = self.ff1(images)
        op2 = self.relu(op1)
        op3 = self.dropout(op2)
        op4 = self.ff2(op3)
        op5 = self.sigmoid(op4)
        return torch.round(op5, decimals=4)

As you can see, it is exactly the same as the solution code. However, I'm still getting an error for this.