How extract loss along with the predicted label from transfomrer model

HI
how I can extract loss along with the predicted label from trained transformer based model?
for example in pytorch we can extarct logits from tokenized inputs:

with torch.no_grad():
output = model(**inputs)
pred_label = torch.argmax(output.logits, axis=-1)
loss = cross_entropy(output.logits, batch[“label”].to(device), reduction=“none”)

how we can do that in TF2?
thanks in advantage