Checklist Before Starting

• Searched for similar PR(s).
• Checked PR Title format
  • In format of: [modules] type: Title
  • modules are in fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data
  • type is in feat, fix, refactor, chore
  • can involve multiple modules, seperated by , or space, like [megatron, fsdp, doc] feat: xxx

What does this PR do?

Fix a TensorDict usage problem which happens at the begin of the training. TensorDict does not support membership checks with the in keyword.

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluatuion results, etc.

High-Level Design

Demonstrate the high-level design if this PR is complex.

Specific Changes

List the specific changes.

API

Demonstrate how the API changes if any.

Usage Example

Provide usage example(s) for easier usage.

# Add code snippet or script demonstrating how to use this 

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks.
• Add [BREAKING] to the PR title description if it breaks any API.
• Update the documentation about your changes in the docs.
• New CI unit test(s) are added to cover the code path.
• Rely on existing unit tests on CI that covers the code path.