get filenames from local dir instead of remote

Signed-off-by: Ryan Nazareth <[email protected]>
flyteorg · Dec 5, 2022 · 23c8bea · 23c8bea
1 parent d6ae50c
commit 23c8bea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flytekit/extras/tensorflow/record.py b/flytekit/extras/tensorflow/record.py
@@ -162,9 +162,9 @@ def to_python_value(
 
         uri, metadata = extract_metadata_and_uri(lv, expected_python_type)
         local_dir = ctx.file_access.get_random_local_directory()
-        files = os.scandir(uri)
-        filenames = [os.path.join(local_dir, f.name) for f in files]
         ctx.file_access.get_data(uri, local_dir, is_multipart=True)
+        files = os.scandir(local_dir)
+        filenames = [os.path.join(local_dir, f.name) for f in files]
         return tf.data.TFRecordDataset(
             filenames=filenames,
             compression_type=metadata.compression_type,