Hi @abercher, regarding your questions:

You can certainly use Transformers as feature extractors on a CPU - since the weights are frozen, you only need the forward pass which is relatively quick to compute
My experience has generally been that you can get significantly worse results when using the last hidden states as features vs fine-tuning end-to-end (in some cases > 20 F1 points!). But this varies depending on the dataset / task, so I am not sure if it’s also true for IMDB.
One thing that seems a bit odd is the hidden.permute(0, 2, 1) part of your forward pass - why do you do this?
You might get better results by using the average of the unmasked hidden state, e.g. in numpy code:

  input_ids = torch.tensor(batch["input_ids"]).to(device)
  attention_mask = torch.tensor(batch["attention_mask"]).to(device)
  with torch.no_grad():
    last_hidden_state = model(input_ids, attention_mask).last_hidden_state
    last_hidden_state = last_hidden_state.cpu().numpy()
  # Use average of unmasked hidden states for classification
  lhs_shape = last_hidden_state.shape
  boolean_mask = ~np.array(batch["attention_mask"]).astype(bool)
  boolean_mask = np.repeat(boolean_mask, lhs_shape[-1], axis=-1)
  boolean_mask = boolean_mask.reshape(lhs_shape)
  masked_mean = np.ma.array(last_hidden_state, mask=boolean_mask).mean(axis=1)
  batch["hidden_state"] = masked_mean.data

You could also see what the performance looks like with something simpler like logistic regression, e.g. first extract the features:
def forward_pass(batch):
  input_ids = torch.tensor(batch["input_ids"]).to(device)
  attention_mask = torch.tensor(batch["attention_mask"]).to(device)
  with torch.no_grad():
    last_hidden_state = model(input_ids, attention_mask, return_dict=True).last_hidden_state
    last_hidden_state = last_hidden_state.cpu().numpy()
  # Use average of unmasked hidden states for classification
  lhs_shape = last_hidden_state.shape
  boolean_mask = ~np.array(batch["attention_mask"]).astype(bool)
  boolean_mask = np.repeat(boolean_mask, lhs_shape[-1], axis=-1)
  boolean_mask = boolean_mask.reshape(lhs_shape)
  masked_mean = np.ma.array(last_hidden_state, mask=boolean_mask).mean(axis=1)
  batch["hidden_state"] = masked_mean.data
  return batch

and then extract the hidden states from your tokenized dataset, e.g. imdb_enc:
imdb_enc = imdb_enc.map(forward_pass, batched=True, batch_size=16)
X_train = np.array(imdb_enc["train"]["hidden_state"])
X_test = np.array(imdb_enc["test"]["hidden_state"])
y_train = np.array(imdb_enc["train"]["label"])
y_test= np.array(imdb_enc["test"]["label"])

and then train a classifier
from sklearn.linear_model import LogisticRegression
lr_clf = LogisticRegression(n_jobs=-1, penalty="none")
lr_clf.fit(X_train, y_train)
lr_clf.score(X_valid, y_valid)

If the result is still bad then you might have to try something more elaborate like averaging over certain layers as was done in the BERT paper.
HTH!