Update pages/13_FFNN.py

pages/13_FFNN.py

@@ -26,7 +26,7 @@ class FeedforwardNeuralNetwork(nn.Module):
         return x
 
 # Function to load the data
-@st.cache
+@st.cache_data
 def load_data():
     transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
     trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
@@ -68,17 +68,15 @@ trainloader, testloader = load_data()
 
 # Streamlit sidebar for input parameters
 st.sidebar.header('Model Hyperparameters')
-input_size = st.sidebar.slider('Input Size', 784, 784, 784)
 hidden1_size = st.sidebar.slider('Hidden Layer 1 Size', 128, 1024, 512)
 hidden2_size = st.sidebar.slider('Hidden Layer 2 Size', 128, 1024, 256)
 hidden3_size = st.sidebar.slider('Hidden Layer 3 Size', 128, 1024, 128)
-output_size = st.sidebar.slider('Output Size', 10, 10, 10)
 learning_rate = st.sidebar.slider('Learning Rate', 0.001, 0.1, 0.01, step=0.001)
 momentum = st.sidebar.slider('Momentum', 0.0, 1.0, 0.9, step=0.1)
 epochs = st.sidebar.slider('Epochs', 1, 20, 5)
 
 # Create the network
-net = FeedforwardNeuralNetwork(input_size, hidden1_size, hidden2_size, hidden3_size, output_size)
+net = FeedforwardNeuralNetwork(784, hidden1_size, hidden2_size, hidden3_size, 10)
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=momentum)