diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp
index 7013f732a..62a6dd066 100644
--- a/src/model/vae/vae.hpp
+++ b/src/model/vae/vae.hpp
@@ -156,7 +156,7 @@ struct VAE : public GGMLRunner {
             output = _compute(n_threads, input, false);
         }
 
-        free_compute_buffer();
+        runner_done();
 
         if (output.empty()) {
             LOG_ERROR("vae encode compute failed");
@@ -207,7 +207,7 @@ struct VAE : public GGMLRunner {
             output = _compute(n_threads, input, true);
         }
 
-        free_compute_buffer();
+        runner_done();
 
         if (output.empty()) {
             LOG_ERROR("vae decode compute failed");