1from huggingface_hub import hf_hub_download
2import torch, torch.nn as nn
3
4class ImprovedMLP(nn.Module):
5 def __init__(self):
6 super().__init__()
7 self.net = nn.Sequential(
8 nn.Flatten(),
9 nn.Linear(784,2048), nn.LayerNorm(2048), nn.GELU(), nn.Dropout(0.1),
10 nn.Linear(2048,1024), nn.LayerNorm(1024), nn.GELU(), nn.Dropout(0.1),
11 nn.Linear(1024,512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(0.1),
12 nn.Linear(512,256), nn.LayerNorm(256), nn.GELU(),
13 nn.Linear(256,128), nn.LayerNorm(128), nn.GELU(),
14 nn.Linear(128,10)
15 )
16 def forward(self,x): return self.net(x)
17
18path = hf_hub_download("chandu1617/MNIST_with_MLP", "mlp_best_fold4.pth")
19model = ImprovedMLP()
20model.load_state_dict(torch.load(path))
21model.eval()
