diff --git a/example/heatEquation2D/src/analyticalSolution.hpp b/example/heatEquation2D/src/analyticalSolution.hpp
index 161d37869b3..953a5da8797 100644
--- a/example/heatEquation2D/src/analyticalSolution.hpp
+++ b/example/heatEquation2D/src/analyticalSolution.hpp
@@ -39,14 +39,14 @@ auto validateSolution(
     double maxError = 0.0;
     for(uint32_t j = 1; j < extent[0] - 1; ++j)
     {
-        for(uint32_t i = 0; i < extent[1]; ++i)
+        for(uint32_t i = 1; i < extent[1] - 1; ++i)
         {
             auto const error = std::abs(buffer.data()[j * extent[1] + i] - exactSolution(i * dx, j * dy, tMax));
             maxError = std::max(maxError, error);
         }
     }
 
-    constexpr double errorThreshold = 1e-5;
+    constexpr double errorThreshold = 1e-4;
     return std::make_pair(maxError < errorThreshold, maxError);
 }
 
diff --git a/example/heatEquation2D/src/heatEquation2D.cpp b/example/heatEquation2D/src/heatEquation2D.cpp
index 458a632d2e4..b186e9ec6db 100644
--- a/example/heatEquation2D/src/heatEquation2D.cpp
+++ b/example/heatEquation2D/src/heatEquation2D.cpp
@@ -55,16 +55,17 @@ auto example(TAccTag const&) -> int
     constexpr alpaka::Vec<Dim, Idx> haloSize{2, 2};
     constexpr alpaka::Vec<Dim, Idx> extent = numNodes + haloSize;
 
-    constexpr uint32_t numTimeSteps = 100;
-    constexpr double tMax = 0.001;
+    constexpr uint32_t numTimeSteps = 4000;
+    constexpr double tMax = 0.1;
+
     // x, y in [0, 1], t in [0, tMax]
     constexpr double dx = 1.0 / static_cast<double>(extent[1] - 1);
     constexpr double dy = 1.0 / static_cast<double>(extent[0] - 1);
     constexpr double dt = tMax / static_cast<double>(numTimeSteps);
 
     // Check the stability condition
-    constexpr double r = dt / std::min(dx * dx, dy * dy);
-    if constexpr(r > 0.5)
+    double r = 2 * dt / ((dx * dx * dy * dy) / (dx * dx + dy * dy));
+    if(r > 1.)
     {
         std::cerr << "Stability condition check failed: dt/min(dx^2,dy^2) = " << r
                   << ", it is required to be <= 0.5\n";
@@ -169,6 +170,7 @@ auto example(TAccTag const&) -> int
 #ifdef PNGWRITER_ENABLED
         if((step - 1) % 100 == 0)
         {
+            alpaka::wait(computeQueue);
             alpaka::memcpy(dumpQueue, uBufHost, uCurrBufAcc);
             alpaka::wait(dumpQueue);
             writeImage(step - 1, uBufHost);
diff --git a/example/heatEquation2D/src/writeImage.hpp b/example/heatEquation2D/src/writeImage.hpp
index dc35867aa48..4b8b6fc56c7 100644
--- a/example/heatEquation2D/src/writeImage.hpp
+++ b/example/heatEquation2D/src/writeImage.hpp
@@ -8,6 +8,7 @@
 
 #include <pngwriter.h>
 
+#include <cmath>
 #include <cstdint>
 #include <iomanip>
 #include <sstream>
@@ -31,7 +32,12 @@ auto writeImage(uint32_t const currentStep, T_Buffer const& buffer) -> void
         for(uint32_t x = 0; x < extents[1]; ++x)
         {
             auto p = buffer.data()[y * extents[1] + x];
-            png.plot(x + 1, extents[0] - y, p, 0., 1. - p);
+            png.plot(
+                x + 1,
+                extents[0] - y,
+                2 * std::exp(std::sqrt(p)) / std::exp(std::sqrt(2)) - 1,
+                0.4,
+                2 - 2 * std::exp(std::sqrt(p)) / std::exp(std::sqrt(2)));
         }
     }
     png.close();