diff --git a/source/leisaac/leisaac/enhance/datasets/hdf5_dataset_file_handler.py b/source/leisaac/leisaac/enhance/datasets/hdf5_dataset_file_handler.py
index dd793eb1..410a95bf 100644
--- a/source/leisaac/leisaac/enhance/datasets/hdf5_dataset_file_handler.py
+++ b/source/leisaac/leisaac/enhance/datasets/hdf5_dataset_file_handler.py
@@ -62,7 +62,7 @@ def _do_write_episode(self, h5_episode_group: h5py.Group, episode: EpisodeData):
             def create_dataset_helper(group, key, value):
                 """Helper method to create dataset that contains recursive dict objects."""
                 if isinstance(value, dict):
-                    key_group = group.get(key, group.create_group(key))
+                    key_group = group[key] if key in group else group.create_group(key)
                     for sub_key, sub_value in value.items():
                         create_dataset_helper(key_group, sub_key, sub_value)
                 else:
@@ -112,7 +112,12 @@ def write_episode(self, episode: EpisodeData, write_mode: StreamWriteMode):
             return
 
         group_name = f"demo_{self._demo_count}"
-        h5_episode_group = self._hdf5_data_group.get(group_name, self._hdf5_data_group.create_group(group_name))
+        h5_episode_group = (
+            self._hdf5_data_group[group_name]
+            if group_name in self._hdf5_data_group
+            else self._hdf5_data_group.create_group(group_name)
+        )
+
         # store number of steps taken
         if "actions" in episode.data:
             if "num_samples" not in h5_episode_group.attrs: