[Feature] Enable storing rollouts on a different device

[Xmaster6y]

Description

Added a storing_device argument to env.rollout to save to a different device, like in collectors.

Motivation and Context

Fix #3177

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• New feature (non-breaking change which adds core functionality)

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!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3199

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 658b219 with merge base 80bfa6e ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

I like it!
Let's just make sure we don't incur any overhead for other use cases

[Xmaster6y]

What overhead? .to(None) does nothing, same for same device tono?

[vmoens]

What overhead? .to(None) does nothing, same for same device tono?

A function call is a function call, I'd prefer if we were avoiding any call to to if not necessary (for clarity)

[Xmaster6y]

Gotcha!

[vmoens]

A small nit otherwise lgtm!