diff --git a/src/UQpy/stochastic_process/KarhunenLoeveExpansion2D.py b/src/UQpy/stochastic_process/KarhunenLoeveExpansion2D.py
index 26a075159..f6fa146ad 100644
--- a/src/UQpy/stochastic_process/KarhunenLoeveExpansion2D.py
+++ b/src/UQpy/stochastic_process/KarhunenLoeveExpansion2D.py
@@ -54,8 +54,11 @@ def __init__(
             self.run(n_samples=self.n_samples, random_variables=random_variables)
 
     def _precompute_one_dimensional_correlation_function(self):
-        self.quasi_correlation_function = np.diagonal(self.correlation_function, axis1=0, axis2=1)
-        self.quasi_correlation_function = np.einsum('ij... -> ...ij', self.quasi_correlation_function)
+        self.quasi_correlation_function = np.zeros(
+            [self.correlation_function.shape[1], self.correlation_function.shape[2],
+             self.correlation_function.shape[3]])
+        for i in range(self.correlation_function.shape[0]):
+            self.quasi_correlation_function[i] = self.correlation_function[i, i]
         self.w, self.v = np.linalg.eig(self.quasi_correlation_function)
         if np.linalg.norm(np.imag(self.w)) > 0:
             print('Complex in the eigenvalues, check the positive definiteness')
@@ -76,7 +79,7 @@ def run(self, n_samples, random_variables=None):
         class. If `n_samples` is provided when the :class:`.KarhunenLoeveExpansion2D` object is defined, the :meth:`run`
         method is automatically called. The user may also call the :meth:`run` method directly to generate samples.
         The :meth:`run`` method of the :class:`.KarhunenLoeveExpansion2D` class can be invoked many times and each time
-        the generated samples areappended to the existing samples.
+        the generated samples are appended to the existing samples.
 
         :param n_samples: Number of samples of the stochastic process to be simulated.
          If the :meth:`run` method is invoked multiple times, the newly generated samples will be appended to the
@@ -92,8 +95,6 @@ def run(self, n_samples, random_variables=None):
             assert (random_variables.shape == (self.thresholds[1], self.thresholds[0], n_samples))
         for i in range(self.one_dimensional_correlation_function.shape[0]):
             if self.thresholds is not None:
-                print(self.w.shape)
-                print(self.v.shape)
                 samples += np.einsum('x, xt, nx -> nxt', np.sqrt(self.w[:, i]), self.v[:, :, i],
                                      KarhunenLoeveExpansion(n_samples=n_samples,
                                                             correlation_function=
diff --git a/tests/unit_tests/stochastic_process/test_karhunen_loeve_2d.py b/tests/unit_tests/stochastic_process/test_karhunen_loeve_2d.py
index ccc825e2f..1bd2fcd21 100644
--- a/tests/unit_tests/stochastic_process/test_karhunen_loeve_2d.py
+++ b/tests/unit_tests/stochastic_process/test_karhunen_loeve_2d.py
@@ -1,7 +1,7 @@
 from UQpy.stochastic_process.KarhunenLoeveExpansion2D import KarhunenLoeveExpansion2D
 import numpy as np
 
-n_samples = 100  # Num of samples
+n_samples = 100_000  # Num of samples
 nx, nt = 20, 10
 dx, dt = 0.05, 0.1
 
@@ -13,7 +13,7 @@
 # R(x_1, x_2, t_1, t_2) = exp(-(x_1 - x_2) ** 2 -(t_1 - t_2) ** 2)
 
 KLE_Object = KarhunenLoeveExpansion2D(n_samples=n_samples, correlation_function=R,
-                                                time_intervals=np.array([dt, dx]), thresholds=[4, 5], random_state=128)
+                                      time_intervals=np.array([dt, dx]), thresholds=[4, 5], random_state=128)
 samples = KLE_Object.samples
 
 
@@ -22,11 +22,16 @@ def test_samples_shape():
 
 
 def test_samples_values():
-    assert np.isclose(samples[13, 0, 13, 6], -1.4385717324103362, rtol=0.01)
+    assert np.isclose(samples[13, 0, 13, 6], -1.8616264088843137, rtol=0.01)
 
 
 def test_run_method():
-    nsamples_second_run = 100
+    nsamples_second_run = 100_000
     KLE_Object.run(n_samples=nsamples_second_run)
     samples = KLE_Object.samples
-    assert samples.shape == (n_samples + nsamples_second_run, 1, len(x), len(t))
\ No newline at end of file
+    assert samples.shape == (n_samples + nsamples_second_run, 1, len(x), len(t))
+
+
+def test_empirical_correlation():
+    assert np.isclose(np.mean(samples[:, 0, 5, 6] * samples[:, 0, 3, 2]), R[5, 3, 6, 2], rtol=0.01)
+    assert np.isclose(np.mean(samples[:, 0, 13, 3] * samples[:, 0, 8, 4]), R[13, 8, 3, 4], rtol=0.01)