Repair parameter check problem in TFRecordDataset

This commit is contained in:
ms_yan 2020-03-31 18:46:59 +08:00 committed by chang zherui
parent bd4a206ae8
commit 51feea03a4
1 changed files with 5 additions and 0 deletions

View File

@ -398,6 +398,7 @@ def check_tfrecorddataset(method):
nreq_param_int = ['num_samples', 'num_parallel_workers', 'num_shards', 'shard_id']
nreq_param_list = ['columns_list']
nreq_param_bool = ['shard_equal_rows']
# check dataset_files; required argument
dataset_files = param_dict.get('dataset_files')
@ -410,6 +411,10 @@ def check_tfrecorddataset(method):
check_param_type(nreq_param_list, param_dict, list)
check_param_type(nreq_param_bool, param_dict, bool)
check_sampler_shuffle_shard_options(param_dict)
return method(*args, **kwargs)
return new_method