[Feature] Implemented device argument for modules.models

[yushiyangk]

Description

Implemented a device argument for modules in torchrl.modules.models to allow them to be directly constructed on a particular device, rather than first constructing it then moving it with .to(). This follows the behaviour of most modules in torch.nn. Updated test_modules.py to make use of this.

However, an arbitrary nn.Module subclass does not necessarily support the device argument, in particular for modules that do not contain parameters or buffers. The utility function added in torchrl/modules/models/utils.py addresses this. This functionality is tested in test_mlp in test_modules.py, parameterised with both nn.ReLU (which does not support device) and nn.PReLU (which does).

Motivation and Context

Addresses #507

Types of changes

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

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

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.

[vmoens]

try / except are a bit expensive
Even if we don't create modules very often in training loops, we do in the tests, and on the long run this risks to have a negative effect on our test speed.
Do you think we can find some other way of doing this?

[codecov]

Codecov Report

Merging #524 (e9cee79) into main (f351fd7) will decrease coverage by 0.17%.
The diff coverage is n/a.

❗ Current head e9cee79 differs from pull request most recent head 3bfc354. Consider uploading reports for the commit 3bfc354 to get more accurate results

@@            Coverage Diff             @@
##             main     #524      +/-   ##
==========================================
- Coverage   86.72%   86.55%   -0.18%     
==========================================
  Files         121      121              
  Lines       21671    21622      -49     
==========================================
- Hits        18795    18715      -80     
- Misses       2876     2907      +31     

Flag	Coverage Δ
linux-cpu	85.08% <0.00%> (-0.02%)	⬇️
linux-gpu	86.39% <0.00%> (-0.16%)	⬇️
linux-outdeps-gpu	77.74% <0.00%> (-0.11%)	⬇️
linux-stable-cpu	85.06% <0.00%> (-0.02%)	⬇️
linux-stable-gpu	86.39% <0.00%> (-0.16%)	⬇️
macos-cpu	84.85% <0.00%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
test/test_collector.py	90.18% <0.00%> (-6.48%)	⬇️
torchrl/trainers/loggers/mlflow.py	78.68% <0.00%> (-3.28%)	⬇️
torchrl/trainers/loggers/wandb.py	80.00% <0.00%> (-2.67%)	⬇️
torchrl/collectors/collectors.py	67.26% <0.00%> (-1.70%)	⬇️
test/test_helpers.py	88.62% <0.00%> (-0.34%)	⬇️
test/test_loggers.py	98.83% <0.00%> (-0.01%)	⬇️
test/test_cost.py	96.41% <0.00%> (ø)
test/test_tensordictmodules.py	98.42% <0.00%> (ø)
test/test_transforms.py	95.79% <0.00%> (+0.01%)	⬆️
torchrl/data/tensordict/tensordict.py	82.56% <0.00%> (+0.09%)	⬆️
... and 2 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[vmoens]

brilliant!

[vmoens]

This is reaaaally cool!
Thanks a million!