[DO NOT MERGE] 2.8 RC Test

.ci/docker/requirements.txt

@@ -38,8 +38,8 @@ tensorboard
 jinja2==3.1.3
 pytorch-lightning
 torchx
-torchrl==0.7.2
-tensordict==0.7.2
+torchrl==0.9.2
+tensordict==0.9.1
 # For ax_multiobjective_nas_tutorial.py
 ax-platform>=0.4.0,<0.5.0
 nbformat>=5.9.2

.jenkins/build.sh

@@ -22,10 +22,13 @@ sudo apt-get install -y pandoc
 #Install PyTorch Nightly for test.
 # Nightly - pip install --pre torch torchvision torchaudio -f https://download.pytorch.org/whl/nightly/cu102/torch_nightly.html
 # Install 2.5 to merge all 2.4 PRs - uncomment to install nightly binaries (update the version as needed).
-# sudo pip uninstall -y fbgemm-gpu torchrec
-# sudo pip uninstall -y torch torchvision torchaudio torchtext torchdata torchrl tensordict
-# sudo pip3 install fbgemm-gpu==1.1.0 torchrec==1.0.0 --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu124
-# pip3 install torch==2.7.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu126
+#sudo pip uninstall -y fbgemm-gpu torchrec
+sudo pip uninstall -y torch torchvision torchaudio torchtext torchdata torchrl tensordict
+pip3 install torch==2.8.0 torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu126
+pip3 install torchrl==0.9.1 tensordict==0.9.1
+# sudo pip3 install fbgemm-gpu==1.2.0 torchrec==1.2.0 --no-cache-dir --index-url https://download.pytorch.org/whl/test/cu126
+
+
 # Install two language tokenizers for Translation with TorchText tutorial
 python -m spacy download en_core_web_sm
 python -m spacy download de_core_news_sm

intermediate_source/dqn_with_rnn_tutorial.py

@@ -342,7 +342,9 @@
 # will return a new instance of the LSTM (with shared weights) that will
 # assume that the input data is sequential in nature.
 #
-policy = Seq(feature, lstm.set_recurrent_mode(True), mlp, qval)
+from torchrl.modules import set_recurrent_mode
+
+policy = Seq(feature, lstm, mlp, qval)
 
 ######################################################################
 # Because we still have a couple of uninitialized parameters we should
@@ -389,7 +391,10 @@
 #   For the sake of efficiency, we're only running a few thousands iterations
 #   here. In a real setting, the total number of frames should be set to 1M.
 #
-collector = SyncDataCollector(env, stoch_policy, frames_per_batch=50, total_frames=200, device=device)
+
+collector = SyncDataCollector(
+    env, stoch_policy, frames_per_batch=50, total_frames=200, device=device
+)
 rb = TensorDictReplayBuffer(
     storage=LazyMemmapStorage(20_000), batch_size=4, prefetch=10
 )
@@ -422,7 +427,8 @@
     rb.extend(data.unsqueeze(0).to_tensordict().cpu())
     for _ in range(utd):
         s = rb.sample().to(device, non_blocking=True)
-        loss_vals = loss_fn(s)
+        with set_recurrent_mode(True):
+            loss_vals = loss_fn(s)
         loss_vals["loss"].backward()
         optim.step()
         optim.zero_grad()
@@ -464,5 +470,5 @@
 #
 # Further Reading
 # ---------------
-# 
+#
 # - The TorchRL documentation can be found `here <https://pytorch.org/rl/>`_.