diff --git a/examples/analytical-function/example_analytical_function.py b/examples/analytical-function/example_analytical_function.py
index 1d0c90160ea4901818aadf228e0a71592ee68e59..8d4a6a064e0dbc0f573ed49c88746ecdbcb0e213 100644
--- a/examples/analytical-function/example_analytical_function.py
+++ b/examples/analytical-function/example_analytical_function.py
@@ -32,7 +32,7 @@ from bayesvalidrox import PyLinkForwardModel, Input, ExpDesigns, PCE, PostProces
 if __name__ == "__main__":
 
     # Number of parameters
-    ndim = 10  # 2, 10
+    ndim = 2  # 2, 10
 
     # =====================================================
     # =============   COMPUTATIONAL MODEL  ================
@@ -212,8 +212,8 @@ if __name__ == "__main__":
     
     # Run using the engine
     engine = Engine(MetaModelOpts, Model, ExpDesign)
-    engine.train_sequential()
-    #engine.train_normal()
+    #engine.train_sequential()
+    engine.train_normal()
 
     # Load the objects
     # with open(f"PCEModel_{Model.name}.pkl", "rb") as input:
@@ -231,9 +231,10 @@ if __name__ == "__main__":
     PostPCE.check_accuracy(n_samples=300)
 
     # Compute the moments and compare with the Monte-Carlo reference
-    # TODO: generalize the moment calculation
-    #if MetaModelOpts.meta_model_type != 'GPE':
     PostPCE.plot_moments()
+    
+    # Plot 3D output of the MetaModel
+    PostPCE.plot_metamodel_3d()
 
     # Plot the evolution of the KLD,BME, and Modified LOOCV error
     #if MetaModelOpts.ExpDesign.method == 'sequential':
@@ -241,8 +242,6 @@ if __name__ == "__main__":
     #    PostPCE.plot_seq_design_diagnostics(refBME_KLD)
 
     # Plot the sobol indices
-    # TODO: generalize this by using moment calculation(?)
-    #if MetaModelOpts.meta_model_type != 'GPE':
     total_sobol = PostPCE.sobol_indices()
 
     # =====================================================