ljt019's picture
dev: uploaded model files
31f46ed verified
raw
history blame contribute delete
327 Bytes
import torch.nn as nn
import torch
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(784, 128)
self.fc2 = nn.Linear(128, 2)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x