Changed model architecture