hey @bennicholl you might have better luck using the Trainer class to run your training with a custom loss function (see docs):
import torch
from transformers import Trainer

class MultilabelTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        loss_fct = torch.nn.BCEWithLogitsLoss()
        loss = loss_fct(logits.view(-1, self.model.config.num_labels),
                        labels.float().view(-1, self.model.config.num_labels))
        return (loss, outputs) if return_outputs else loss

then you can define a compute_metrics function as follows
from scipy.special import expit as sigmoid
from sklearn.metrics import classification_report

def compute_metrics(pred):
    y_true = pred.label_ids
    y_pred = sigmoid(pred.predictions) y_pred = (y_pred>0.5).astype(float)
        
    clf_dict = classification_report(y_true, y_pred, target_names=all_labels,
                                         zero_division=0, output_dict=True)
    return {"micro f1": clf_dict['micro avg']['f1-score'], "macro f1": clf_dict['macro avg']['f1-score']}

and after fine-tuning you can then get the predictions on your validation set via Trainer.predict:
trainer = # fine-tuned Trainer
pred = trainer.predict(your_eval_dataset)
metrics = compute_metrics(pred)

hth!