diff --git a/examples/save-load-state/save-load-state.cpp b/examples/save-load-state/save-load-state.cpp
index 61c71c358..b1856e592 100644
--- a/examples/save-load-state/save-load-state.cpp
+++ b/examples/save-load-state/save-load-state.cpp
@@ -1,5 +1,6 @@
 #include "common.h"
 #include "llama.h"
+#include "llama-util.h"
 #include "build-info.h"
 
 #include <vector>
@@ -66,10 +67,9 @@ int main(int argc, char ** argv) {
 
     // Save state (rng, logits, embedding and kv_cache) to file
     {
-        FILE *fp_write = fopen("dump_state.bin", "wb");
-        llama_copy_state_data(ctx, state_mem); // could also copy directly to memory mapped file
-        fwrite(state_mem, 1, state_size, fp_write);
-        fclose(fp_write);
+        llama_file file("dump_state.bin", "wb");
+        llama_data_file_context data_ctx(&file);
+        llama_copy_state_data(ctx, &data_ctx); // can also copy to memory buffer
     }
 
     // save state (last tokens)