diff --git a/otbtf/tfrecords.py b/otbtf/tfrecords.py
index 7d1a4d4c..254efd2a 100644
--- a/otbtf/tfrecords.py
+++ b/otbtf/tfrecords.py
@@ -172,8 +172,8 @@ def read(self, batch_size, target_keys, n_workers=1, drop_remainder=True, shuffl
         :param num_parallel_calls: number of parallel calls for the parsing + preprocessing step
         :param kwargs: some keywords arguments for preprocessing_fn
         """
-        for file in [self.output_types_file, self.output_shapes_file]:
-            assert self.output_types, f"The file {file} is missing!"
+        for dic, file in zip([self.output_types, self.output_shapes], [self.output_types_file, self.output_shapes_file]):
+            assert dic, f"The file {file} is missing!"
 
         options = tf.data.Options()
         if shuffle_buffer_size: