nn.Linear(model.config.hidden_size, 512),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.Dropout(.5),
nn.Linear(512, 256),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(.5),
nn.Linear(256, 2)