add comment

EleutherAI · Apr 27, 2023 · 72fc632 · 72fc632
1 parent 2cd4401
commit 72fc632
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/elk/utils/fsdp.py b/elk/utils/fsdp.py
@@ -165,6 +165,8 @@ def imap(
  closure_pkl = dill.dumps(closure)
  result_queues = []
  for q, result_queue in zip(self._task_queues, self._result_queues):
+ # Put the same dataset on each queue, so that each worker gets the same
+ # inputs
  q.put((closure_pkl, dataset))
  result_queues.append(result_queue)
 
@@ -213,7 +215,7 @@ def _worker(
  logging.disable(logging.CRITICAL)
  warnings.filterwarnings("ignore")
 
- closure: Callable[[ModelOutput], Any] | None = None
+ closure: Callable[[ModelOutput], Any]
  dataset: Dataset | None = None
  device = devices[rank]