I found this 31 and changed it slightly to adapt it to your question. Haven’t done any rigorous testing so you should probably do that.
def F1_score(prob, label):
    prob = prob.bool()
    label = label.bool()
    epsilon = 1e-7
    TP = (prob & label).sum().float()
    TN = ((~prob) & (~label)).sum().float()
    FP = (prob & (~label)).sum().float()
    FN = ((~prob) & label).sum().float()
    #accuracy = (TP+TN)/(TP+TN+FP+FN)
    precision = torch.mean(TP / (TP + FP + epsilon))
    recall = torch.mean(TP / (TP + FN + epsilon))
    F2 = 2 * precision * recall / (precision + recall + epsilon)
    return precision, recall, F2

y_true = torch.tensor([[1,0,0,1]])
y_pred = torch.tensor([[1,1,0,0]])
print(F1_score(y_pred, y_true))

Edit:
Found that PyTorch Lightning does have implementation for precision-recall 67 and f1 score 65, perhaps you could use that.