diff --git a/src/pao_model.F b/src/pao_model.F
index 73cbd3162d..04e16f5819 100644
--- a/src/pao_model.F
+++ b/src/pao_model.F
@@ -361,6 +361,7 @@ CONTAINS
       CPASSERT(SIZE(predicted_xblock, 1) == n)
       CPASSERT(SIZE(predicted_xblock, 2) == m)
       CPASSERT(SIZE(predicted_xblock, 3) == 1)
+      CPASSERT(ALL(predicted_xblock == predicted_xblock)) ! checking for NaNs
       IF (PRESENT(block_X)) THEN
          block_X = RESHAPE(predicted_xblock, [n*m, 1])
       END IF
@@ -375,6 +376,7 @@ CONTAINS
          NULLIFY (edge_vectors_grad)
          CALL torch_tensor_data_ptr(edge_vectors_grad_tensor, edge_vectors_grad)
          CPASSERT(SIZE(edge_vectors_grad, 1) == 3 .AND. SIZE(edge_vectors_grad, 2) == num_edges)
+         CPASSERT(ALL(edge_vectors_grad == edge_vectors_grad)) ! checking for NaNs
          DO iedge = 1, num_edges
             jneighbor = INT(edge_index(iedge, 1) + 1)
             kneighbor = INT(edge_index(iedge, 2) + 1)
diff --git a/tools/pao-ml/pao/model.py b/tools/pao-ml/pao/model.py
index 8d75f50168..c556d5dd1c 100644
--- a/tools/pao-ml/pao/model.py
+++ b/tools/pao-ml/pao/model.py
@@ -75,6 +75,8 @@ class PaoModel(SequentialGraphNetwork):  # type: ignore
         prim_basis_specs = {
             "DZVP-MOLOPT-GTH/H": "2x0e + 1x1o",  # two s-shells, one p-shell
             "DZVP-MOLOPT-GTH/O": "2x0e + 2x1o + 1x2e",  # two s, two p, one d-shell
+            "DZVP-MOLOPT-GGA-GTH-q1/H": "2x0e + 1x1o",
+            "DZVP-MOLOPT-GGA-GTH-q6/O": "2x0e + 2x1o + 1x2e",
             "TZV2P-MOLOPT-GGA-GTH-q1/H": "3x0e + 2x1o + 1x2e",
             "TZV2P-MOLOPT-GGA-GTH-q6/O": "3x0e + 3x1o + 2x2e + 1x3o",
         }
@@ -227,7 +229,8 @@ class SymmetricMatrix(torch.nn.Module):
         assert vector.shape[-1] == sum(dim(l) for l in self.irreps_in_ls)
         basis_size = sum(dim(l) for l in self.irreps_prim_basis_ls)
         matrix = torch.zeros(vector.shape[:-1] + (basis_size, basis_size))
-        matrix[..., :, :] = torch.eye(basis_size)  # ensure matrix is diagonalizable
+        # Ensure matrix has distinct eigenvalues as needed for grad of torch.linalg.eigh().
+        matrix[..., :, :] = torch.diag(torch.arange(1, basis_size + 1))
         c = 0  # position in vector
         z = 0  # position in self.wigner_blocks
         for i, li in enumerate(self.irreps_prim_basis_ls):