Revert "try out not sending the sentinel"

This reverts commit d24ded8.

elk/utils/fsdp.py

@@ -248,7 +248,7 @@ def _worker(
  out_queue.put(output_applied)
 
  # Indicate we're done with this dataset
- # out_queue.put_nowait(None)
+ out_queue.put_nowait(None)
 
  # Clean up the FSDP process group
  if fsdp_port is not None:

tests/test_inference_server.py

@@ -36,7 +36,7 @@ def test_inference_server_normal():
  )
 
 
-def test_inference_server_fsdp_single():
+def test_inference_server_fsdp():
  model_str = "sshleifer/tiny-gpt2"
  server = InferenceServer(
  model_str=model_str, num_workers=2, fsdp=True, cpu_offload=True
@@ -51,7 +51,7 @@ def test_inference_server_fsdp_single():
  inputs = dict(input_ids=input_ids)
  # make the dict a dataset, while still making it a pytorch dataset
  input_dataset = Dataset.from_dict(inputs)
- input_dataset.set_format(tyxpe="torch")
+ input_dataset.set_format(type="torch")
  outputs = server.map(dataset=input_dataset, closure=lambda x: x)[0]
  assert (
  type(outputs) == transformers.modeling_outputs.CausalLMOutputWithCrossAttentions