[Feature] Make benchmarked losses compatible with torch.compile

[vmoens]

Stack from ghstack (oldest at bottom):

• [Feature] RB compability with compile #2426
• [CI] Add benchmarks to test runs #2410
• [Doc] Document losses in README.md #2408
• -> [Feature] Make benchmarked losses compatible with torch.compile #2405

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 9 New Failures, 1 Cancelled Job, 10 Unrelated Failures

As of commit e4441a0 with merge base 605b4aa ():

NEW FAILURES - The following jobs have failed:

• Habitat Tests on Linux / tests (3.9, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t b92e5149785b589c2b9c12e5a1266618338375a7685dc53cdcffa7d6828387bb /exec failed with exit code 134
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  test/test_env.py::TestInfoDict::test_auto_register[device0]
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  test/test_env.py::TestInfoDict::test_auto_register[device0]
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  test/test_env.py::TestInfoDict::test_auto_register[device0]
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
  test/test_env.py::TestInfoDict::test_auto_register[device0]
• Unit-tests on Linux / tests-cpu-oldget (3.12) / linux-job (gh)
  test/test_env.py::TestInfoDict::test_auto_register[device0]
• Unit-tests on Linux / tests-gpu (3.11, 12.1) / linux-job (gh)
  test/test_env.py::TestInfoDict::test_auto_register[device0]
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
  test/test_transforms.py::TestKLRewardTransform::test_kl_lstm
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh)
  test/test_env.py::TestInfoDict::test_auto_register[device0]

CANCELLED JOB - The following job was cancelled. Please retry:

• Wheels / build-wheel-windows (3.10, 3.10.3) (gh)

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh) (detected as infra flaky with no log or failing log classifier)
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh) (detected as infra flaky with no log or failing log classifier)
• Unit-tests on Windows / unittests-cpu / windows-job (gh) (detected as infra flaky with no log or failing log classifier)

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cpu (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cuda11_8 (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cuda12_1 (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cuda12_4 (gh) (trunk failure)
• Wheels / build-wheel-windows (3.11, 3.11) (gh) (trunk failure)
• Wheels / build-wheel-windows (3.12, 3.12) (gh) (trunk failure)
• Wheels / build-wheel-windows (3.9, 3.9) (gh) (trunk failure)

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

[vmoens]

Required merges before landing this:
pytorch/pytorch#135000 (comment)
pytorch/pytorch#135001