diff --git a/docs/ftorch.rst b/docs/ftorch.rst
index 936f26c..a0e4f4b 100644
--- a/docs/ftorch.rst
+++ b/docs/ftorch.rst
@@ -191,7 +191,7 @@ To evaluate Skala, we download the model checkpoint from HuggingFace using the `
 
    .. code-block:: bash
 
-      python prepare_inputs.py --output_dir features --molecule H2 --basis def2-QZVP
+      python prepare_inputs.py --output-dir features --basis def2-QZVP
 
    The script needs the ``skala`` package installed in your Python environment, which can be done via pip:
 
@@ -210,7 +210,7 @@ The output for the H2 molecule with the def2-QZVP basis set should look like thi
 .. code-block:: text
 
    [1] Loading model from skala-1.0.fun
-   [2] Loading features from H2-def2qzvp
+   [2] Loading features from features
     -> Loading coarse_0_atomic_coords
     -> Loading grad
     -> Loading grid_coords
@@ -260,8 +260,8 @@ In the output we can see the computed exchange-correlation energy as well as the
 
 .. code-block:: text
 
-   [1] Loading model from ../../cpp/gauxc_integration/skala-1.0.fun
-   [2] Loading features from H2-def2qzvp
+   [1] Loading model from skala-1.0.fun
+   [2] Loading features from features
     -> Loading coarse_0_atomic_coords
     -> Loading grad
     -> Loading grid_coords
@@ -307,4 +307,4 @@ Summary
 
 In this example, we have demonstrated how to use Skala in Fortran with the FTorch library.
 We have shown how to set up the development environment, build the application using CMake, and run the application to compute exchange-correlation energies and potentials using a Skala model.
-This example serves as a starting point for integrating Skala into your Fortran applications using FTorch, and can be extended to include more complex features and functionalities as needed.
\ No newline at end of file
+This example serves as a starting point for integrating Skala into your Fortran applications using FTorch, and can be extended to include more complex features and functionalities as needed.
diff --git a/examples/cpp/cpp_integration/prepare_inputs.py b/examples/cpp/cpp_integration/prepare_inputs.py
index 9f3369f..26be867 100755
--- a/examples/cpp/cpp_integration/prepare_inputs.py
+++ b/examples/cpp/cpp_integration/prepare_inputs.py
@@ -19,15 +19,21 @@ def main() -> None:
         type=Path,
         help="Output directory for generated feature files.",
     )
+    parser.add_argument(
+        "--basis",
+        default="def2-qzvp",
+        type=str,
+        help="Basis set.",
+    )
     args = parser.parse_args()
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
 
-    molecule = gto.Mole(
+    molecule = gto.M(
         atom="H 0 0 0; H 0 0 1",
-        basis="def2-qzvp",
+        basis=args.basis,
         verbose=0,
-    ).build()
+    )
 
     # Create a set of meta-GGA features for this molecule.
     dm = get_density_matrix(molecule)