Update check_tfrecords to use new dataset load function.

[cfezequiel]

Description

Make check_tfrecords use tfrecorder.load function.

This changes the check_tfrecords signature such that the user no longer has to specify the compression_type (inferred from TFRecord filenames) and can specify the dataset split (TRAIN, VALIDATION, TEST) to determine which dataset to inspect.

Type of change

Please delete options that are not relevant.

• Breaking change (fix/feature that may cause existing functionality to not work as expected)

Checklist

Please delete options that are not relevant.

• My code adheres to the Google Python Style Guide
• I ran make pylint and code is rated 10/10
• I have added tests that prove my fix is effective or that my feature works
• I ran make test and all tests pass
• I ran the tool and verified the change works
• I have adequately commented my code, particularly in hard-to-understand areas
• I have made relevant changes to the documentation, if needed
• My changes generate no new warnings
• I have corrected any misspellings in my code

[mbernico]

Nice, straightforward and simplifies the codebase! lgtm