fix lists (mosaicml#3138)

mvpatel2000 · Mar 22, 2024 · 81547ca · 81547ca
1 parent 8cc44ed
commit 81547ca
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
@@ -296,9 +296,9 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner):
             receiver = dist.get_global_rank() != rank_in_first_replica
 
             # Send list of files to all ranks
-            file_list = [
+            file_list = [[
                 file_name for file_name in sorted(os.listdir(self.destination_path)) if file_name.endswith('.distcp')
-            ]
+            ]]
             dist.broadcast_object_list(file_list, src=rank_in_first_replica, group=replicate_process_group)
             file_list = file_list[0]
             log.debug(f'List of files to broadcast: {file_list}')