A simple mlp mnist model using custom dropout class

mlp_mnist.py

class MNISTModel(torch.nn.Module):
    
    def __init__(self):
        super(MNISTModel, self).__init__()
        self.layer_1 = nn.Linear(28 * 28, 512)
        self.layer_2 = nn.Linear(512, 512)
        self.layer_3 = nn.Linear(512, 10)
        self.dropout = Dropout(.5)
    
    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = self.layer_1(x)
        x = self.layer_2(x)
        x = self.dropout(x)
        output = self.layer_3(x)
        return x