I was able to exactly align the loss values for KL-divergence and cross-entropy. See the code below:
a = torch.randint(10, (2, 5))
b = torch.Tensor([[0., 0., 0., 1., 0.], [0., 0., 1., 0., 0.]])
target = torch.Tensor([3, 2]).type(torch.LongTensor)

print('a', a)
print('b', b)

a_log_soft = F.log_softmax(a.float(), dim=1)
print('a_log_soft', a_log_soft)
b_soft = F.softmax(b, dim=1)
print('b_soft', b_soft)

print('kl_loss batchmean', F.kl_div(a_log_soft, b, reduction='batchmean'))

print('kl_loss_prob batchmean', F.kl_div(a_log_soft, b_soft, reduction='batchmean'))

print('cross_loss mean', F.cross_entropy(a.float(), target))

here, kl_loss batchmean aligns perfectly with cross_loss mean. However, kl_loss_prob batchmean doesn’t align with cross_loss mean. However, in a real scenario if we have our b input as raw logits, kl_loss batchmean is the one that should be used.
Also, make sure to use reduction='batchmean'.
Hence, in my original question all I need to do is change:
loss = F.kl_div(log_probs_flat, target_flat)

to
loss = F.kl_div(log_probs_flat, y_onehot, reduction='batchmean')

and this loss will exactly match with the cross-entropy loss.