[Feature] Add NPU support for SyncDataCollector #3155
Merged
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR introduces native NPU device support to SyncDataCollector, enabling it to seamlessly accept CPU, CUDA, MPS, and NPU devices. With this change, users can easily run state-of-the-art implementations on NPUs without modifying the core workflow.
Motivation and Context
We aim to add Ascend NPU support to TorchRL. This PR is a first step that makes it straightforward to run sota-implementations on NPU devices by simply moving all components including DataCollectors onto NPU. Related discussion can be found here: #3154
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
x
in all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!