class MLPEncoder(nn.Module):
    def __init__(self):
        super(MLPEncoder, self).__init__()

        # TODO: Fix hard coding
        self.model = nn.Sequential(OrderedDict([
          ('layer1', nn.Linear(784, 400)),
          ('relu1', nn.ReLU()),
          ('layer2', nn.Linear(400, 400)),
          ('relu2', nn.ReLU()),
          ('layer3', nn.Linear(400, 200)),
          ('relu3', nn.ReLU()),
          ('layer4', nn.Linear(200, 200)),
          ('relu4', nn.ReLU())
        ]))
        print(self.model)
        
        # add hook to store activatiosns
        self.activations = {}
        def store_activations(model, input, output):
            self.activations[model.__name__] = output

        for name, layer in self.model.named_children():
            layer.__name__ = name
            layer.register_forward_hook(store_activations)

    def forward(self, x):
        return self.model(x.view(-1, 784))

    def get_activations(name):
        return self.activations[name]