This particular form is a batch matrix multiplication if you swap A and B, so you could use torch.bmm or torch.matmul directly.
Einsum will (currently) reduce to bmm, so it’ll be similar in terms of performance, but with a few permutations.
Best regards
Thomas