pytorch_Model_medium

MyModel.py

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # Encoder
        self.block1 = nn.Sequential(
            nn.Linear(128,64),
            nn.Tanh(),
            nn.Linear(64, 32),
            nn.Tanh(),
            nn.Linear(32, 16),
            nn.Tanh(),
            nn.Linear(16, 10),
        )
        
    def forward(self, x):
        out = self.block1(x)
        return out