Here is the full code example:
During Train: forward pass, calculate loss, optimize,
During Test: forwards pass, calculate loss, strip last layer, get features for the last layer output.

# dataloaders_dict = contains dict of both loaders  {'train', 'val'}

# add some layers to resnet50
model_ft = models.resnet34(pretrained=True)
model_ft.fc_1 = nn.Sequential( nn.Dropout(0.1), nn.Linear(1000, 512))
model_ft.fc_2 = nn.Sequential( nn.Dropout(0.1), nn.Linear(512, 512))
model_ft.fc_3 = nn.Sequential( nn.Dropout(0.1), nn.Linear(512, 256))
model_ft.fc_4 = nn.Sequential( nn.Dropout(0.1), nn.Linear(256, 128))
model_ft.fc_5 = nn.Linear(128, 2)

model_ft = model_ft.to(device)

loss_function = nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(model_ft.parameters(), lr=1e-4)

def train_or_test(model, dataloaders_dict, optimizer, loss_function, phase):
  # since = time.time()

  running_loss = 0.0
  running_corrects = 0
  results1 = []

  # some training
  if phase == 'train':
    model.train()
  else: # 'val'/ use as a feature extractor to get last but layer values
    model.eval()

  for inputs, labels in dataloaders_dict[phase]:
    inputs = inputs.to(device)
    labels = labels.to(device)

    optimizer.zero_grad()

    outputs = model(inputs)
    loss = loss_function(outputs, labels)
    _, preds = torch.max(outputs, 1)

    # statistics
    running_loss += loss.item() * inputs.size(0)
    running_corrects += torch.sum(preds == labels.data)

    # backward + optimize only if in training phase
    if phase == 'train':
        loss.backward()
        optimizer.step()
    else: #'val'
      # use as extractor
      feature_extractor = torch.nn.Sequential(*list(model.children())[:-1])
      feature_extractor = feature_extractor.to(device)

      # forward pass input, get last but layer feature values
      feature_tensor = feature_extractor(inputs) # output now has the features corresponding to input x
      feature_arr = feature_tensor.cpu().detach().numpy().flatten()

      results1.append([labels[0].cpu().numpy(), feature_arr])

      
  if phase=='train':
    val_loss = running_loss/2488
    val_acc = running_corrects.double()/2488
    return model, val_loss, val_acc
  else: # 'val'
    val_loss = running_loss/278
    val_acc = running_corrects.double()/278
    return val_loss, val_acc, results1


model_ft, loss, acc = train_or_test(model_ft, dataloaders_dict, optimizer, loss_function, phase='train')
features, loss, acc = train_or_test(model_ft, dataloaders_dict, optimizer, loss_function, phase='val')