# Replacing classification head with a new head
# the new training objective is still binary classification,
# except these parameters will be used to decide if a
# review was positive or negative
model.classifier = nn.Linear(d_model, 1, bias=False).to(device)

# resetting the optimizer to have access to the parameters of the new head
optimizer = optim.Adam(model.parameters(), lr=0.001)