diff --git a/lib/src/BufferedWriter.cpp b/lib/src/BufferedWriter.cpp
index feb004b..5d30abf 100644
--- a/lib/src/BufferedWriter.cpp
+++ b/lib/src/BufferedWriter.cpp
@@ -38,7 +38,8 @@ void BufferedWriter::write_metadata_to_file()
 
             vector<size_t> data_shape = {header_data_type.value_shape};
 
-            create_dataset(dataset_name, data_shape, header_data_type.type, header_data_type.endianness, false);
+            create_dataset(dataset_name, data_shape, header_data_type.type, header_data_type.endianness, false, 
+                metadata_buffer->get_n_images());
 
             H5::AtomType dataset_data_type(H5FormatUtils::get_dataset_data_type(header_data_type.type));
             dataset_data_type.setOrder(H5T_ORDER_LE);
diff --git a/lib/src/H5Writer.cpp b/lib/src/H5Writer.cpp
index 3487cd1..7ac8548 100644
--- a/lib/src/H5Writer.cpp
+++ b/lib/src/H5Writer.cpp
@@ -141,7 +141,7 @@ void H5Writer::write_data(const string& dataset_name, const size_t data_index, c
 }
 
 void H5Writer::create_dataset(const string& dataset_name, const vector<size_t>& data_shape, 
-    const string& data_type, const string& endianness, bool chunked)
+    const string& data_type, const string& endianness, bool chunked, hsize_t dataset_size)
 {
     // Number of dimensions in each data point.
     const size_t data_rank = data_shape.size();
@@ -153,7 +153,7 @@ void H5Writer::create_dataset(const string& dataset_name, const vector<size_t>&
     hsize_t dataset_chunking[dataset_rank];
     
     // This should be equivalent to the total number of frames in this file.
-    dataset_dimension[0] = initial_dataset_size;
+    dataset_dimension[0] = dataset_size;
     // This dataset can be resized without limits. 
     max_dataset_dimension[0] = H5S_UNLIMITED;
     // Chunking is always set to a single data point.
@@ -291,7 +291,7 @@ hsize_t H5Writer::prepare_storage_for_data(const string& dataset_name, const siz
 
     // Create the dataset if we don't have it yet.
     if (datasets.find(dataset_name) == datasets.end()) {
-        create_dataset(dataset_name, data_shape, data_type, endianness, true);
+        create_dataset(dataset_name, data_shape, data_type, endianness, true, initial_dataset_size);
     }
 
     hsize_t current_dataset_size = datasets_current_size.at(dataset_name);
diff --git a/lib/src/H5Writer.hpp b/lib/src/H5Writer.hpp
index e74c6e8..1a8b4b3 100644
--- a/lib/src/H5Writer.hpp
+++ b/lib/src/H5Writer.hpp
@@ -29,7 +29,7 @@ class H5Writer
             const std::string& data_type, const std::string& endianness);
 
         void create_dataset(const std::string& dataset_name, const std::vector<size_t>& data_shape, 
-            const std::string& data_type, const std::string& endianness, bool chunked);
+            const std::string& data_type, const std::string& endianness, bool chunked, hsize_t dataset_size);
         
         size_t get_relative_data_index(const size_t data_index);