Zolisa commited on
Commit
0b57a4b
·
verified ·
1 Parent(s): e323bd4

Upload simple_nn.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. simple_nn.py +16 -0
simple_nn.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+
4
+ class SimpleNN(nn.Module):
5
+ def __init__(self, input_size=784, hidden_size=128, num_classes=10):
6
+ super(SimpleNN, self).__init__()
7
+ self.fc1 = nn.Linear(input_size, hidden_size)
8
+ self.relu = nn.ReLU()
9
+ self.fc2 = nn.Linear(hidden_size, num_classes)
10
+
11
+ def forward(self, x):
12
+ x = x.view(x.size(0), -1)
13
+ x = self.fc1(x)
14
+ x = self.relu(x)
15
+ x = self.fc2(x)
16
+ return x