Would it be also correct to do it that way ?
import torch
from torch.distributions import RelaxedOneHotCategorical

batch_param_obtained_from_a_nn = torch.rand(2, 4)
q_m = RelaxedOneHotCategorical(torch.tensor([5.4]), logits=batch_param_obtained_from_a_nn)
z = q_m.rsample()

# ...

p_m = torch.tensor([0.1, 0.2, 0.3, 0.4])

kl = - torch.mean(torch.sum(torch.log(z + 1e-20) * (torch.log(z + 1e-20) - torch.log(p_m + 1e-20)), dim=1))
z                                                                                                                                                                                                                                                  
tensor([[0.3518, 0.2110, 0.2363, 0.2010],
        [0.2213, 0.2337, 0.2818, 0.2632]])

kl                                                                                                                                                                                                                                                 
tensor(0.3673)