RuntimeError: torch.triangular_solve: Expected b to have at least 2 dimensions, but it has 1 dimensions instead

I converted this function from tensorflow to pytorch

import tensorflow as tf
def gaussian_log_density(samples, means, chol_covars):
    covar_logdet = 2 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_covars)), -1)
    diff = tf.expand_dims(samples - means, -1)
    exp_term = tf.reduce_sum(tf.square(tf.linalg.triangular_solve(chol_covars, diff)), axis=[-2, -1])
    return - 0.5 * (tf.cast(tf.shape(samples)[-1], exp_term.dtype) * np.log(2 * np.pi) + covar_logdet + exp_term)

and I am wondering whether I correctly converted tf.linalg.diag_part in tf.reduce_sum(..., axis=[-2, -1]) in my implementation?

def gaussian_log_density(samples, means, chol_covars):
    covar_logdet = 2 *torch.sum(torch.log(torch.diag_embed(chol_covars, offset=0, dim1=-2, dim2=-1)+1e-10),dim=-1)
    diff = torch.unsqueeze(samples - means, -1)
    exp_term = torch.sum(torch.square(torch.triangular_solve(chol_covars, diff)), (-2, -1))
    return - 0.5 * (samples.size()[-1].type( exp_term.dtype) * torch.log(2 * torch.tensor(np.pi)) + covar_logdet + exp_term)

if __name__=="__main__":
  a = torch.tensor([2.0, 1.1, 1.2])
  b =torch.tensor([1, 2.4, 5.35])
  s= torch.tensor([1.02, 0.9, 1.1])
  print(gaussian_log_density(s, a, b))

since I am getting this error just with testing my code:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-44-ab80d557c449> in <module>()
     10   b =torch.tensor([1, 2.4, 5.35])
     11   s= torch.tensor([1.02, 0.9, 1.1])
---> 12   print(gaussian_log_density(s, a, b))

<ipython-input-44-ab80d557c449> in gaussian_log_density(samples, means, chol_covars)
      3     print(covar_logdet)
      4     diff = torch.unsqueeze(samples - means, -1)
----> 5     exp_term = torch.sum(torch.square(torch.triangular_solve(chol_covars, diff)), (-2, -1))
      6     return - 0.5 * (samples.size()[-1].type( exp_term.dtype) * torch.log(2 * torch.tensor(np.pi)) + covar_logdet + exp_term)
      7 

RuntimeError: torch.triangular_solve: Expected b to have at least 2 dimensions, but it has 1 dimensions instead

Thanks.