import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self, config):
        super(Model, self).__init__()
        self.drop = nn.Dropout(config['dropout'])

        self.fc1 = nn.Linear(784, 2000)
        self.fc2 = nn.Linear(2000, 2000)
        self.fc3 = nn.Linear(2000, 2000)
        self.fc4 = nn.Linear(2000, 2000)
        self.fc5 = nn.Linear(2000, 10)

    def forward(self, x):
        # 784 -> 2000
        x = F.relu(self.drop(self.fc1(x)))
        # 2000 -> 2000
        x = F.relu(self.drop(self.fc2(x)))
        # 2000 -> 2000
        x = F.relu(self.drop(self.fc3(x)))
        # 2000 -> 2000
        x = F.relu(self.drop(self.fc4(x)))
        # 2000 -> 100
        x = self.fc5(x)
        return x
[((8, 20, 8, 49), 'torch.nn.Dropout', 'nn.Dropout', ({(8, 31, 8, 48): "config['dropout']"}, {}), "(config['dropout'])", True, 'import torch.nn as nn\n'), ((10, 19, 10, 39), 'torch.nn.Linear', 'nn.Linear', ({(10, 29, 10, 32): '784', (10, 34, 10, 38): '2000'}, {}), '(784, 2000)', True, 'import torch.nn as nn\n'), ((11, 19, 11, 40), 'torch.nn.Linear', 'nn.Linear', ({(11, 29, 11, 33): '2000', (11, 35, 11, 39): '2000'}, {}), '(2000, 2000)', True, 'import torch.nn as nn\n'), ((12, 19, 12, 40), 'torch.nn.Linear', 'nn.Linear', ({(12, 29, 12, 33): '2000', (12, 35, 12, 39): '2000'}, {}), '(2000, 2000)', True, 'import torch.nn as nn\n'), ((13, 19, 13, 40), 'torch.nn.Linear', 'nn.Linear', ({(13, 29, 13, 33): '2000', (13, 35, 13, 39): '2000'}, {}), '(2000, 2000)', True, 'import torch.nn as nn\n'), ((14, 19, 14, 38), 'torch.nn.Linear', 'nn.Linear', ({(14, 29, 14, 33): '2000', (14, 35, 14, 37): '10'}, {}), '(2000, 10)', True, 'import torch.nn as nn\n')]