def compute_A(m, nacv, hbar=1):
    """
    Compute the non-adiabatic coupling term A that shows up in the off-diagonal
    of the Hamiltonian in the adiabatic basis.

    Args:
            m (torch.Tensor): masses of dimension N_at,
            nacv (torch.Tensor): non-adiabatic coupling vector, of dimension
                    n_ij x N_at x 3. n_ij are the subsets of of N_I and N_J with enough
                    overlap to warrant a matrix element calculation.
            hbar (float): reduced Planck constant. Set to 1 when using atomic units.
    Returns:
            A_ij (torch.Tensor): coupling term A, of dimension n_ij.
    """

    m_reshape = m.reshape(1, -1, 1)
    A_ij = (-hbar ** 2 / m_reshape * nacv).sum((1, 2))

    return A_ij