Apply 1 suggestion(s) to 1 file(s)

remicres · Sep 13, 2022 · b5eb853 · b5eb853
1 parent 0c58c8f
commit b5eb853
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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: