[Test] Fix device in PPO tests

ghstack-source-id: efe21e7
Pull-Request-resolved: #2971

Author: Vincent Moens

test/test_cost.py

@@ -8783,6 +8783,7 @@ def test_ppo(
             value,
             loss_critic_type="l2",
             functional=functional,
+            device=device,
         )
         if composite_action_dist:
             loss_fn.set_keys(
@@ -8883,6 +8884,7 @@ def test_ppo_composite_no_aggregate(
             value,
             loss_critic_type="l2",
             functional=functional,
+            device=device,
         )
         loss_fn.set_keys(
             action=("action", "action1"),
@@ -8943,9 +8945,19 @@ def test_ppo_state_dict(
             device=device, composite_action_dist=composite_action_dist
         )
         value = self._create_mock_value(device=device)
-        loss_fn = loss_class(actor, value, loss_critic_type="l2")
+        loss_fn = loss_class(
+            actor,
+            value,
+            loss_critic_type="l2",
+            device=device,
+        )
         sd = loss_fn.state_dict()
-        loss_fn2 = loss_class(actor, value, loss_critic_type="l2")
+        loss_fn2 = loss_class(
+            actor,
+            value,
+            loss_critic_type="l2",
+            device=device,
+        )
         loss_fn2.load_state_dict(sd)
 
     @pytest.mark.parametrize("loss_class", (PPOLoss, ClipPPOLoss, KLPENPPOLoss))
@@ -8993,6 +9005,7 @@ def test_ppo_shared(self, loss_class, device, advantage, composite_action_dist):
             value,
             loss_critic_type="l2",
             separate_losses=True,
+            device=device,
         )
 
         if advantage is not None:
@@ -9100,6 +9113,7 @@ def test_ppo_shared_seq(
             loss_critic_type="l2",
             separate_losses=separate_losses,
             entropy_coef=0.0,
+            device=device,
         )
 
         loss_fn2 = loss_class(
@@ -9108,6 +9122,7 @@ def test_ppo_shared_seq(
             loss_critic_type="l2",
             separate_losses=separate_losses,
             entropy_coef=0.0,
+            device=device,
         )
 
         if advantage is not None:
@@ -9202,7 +9217,12 @@ def test_ppo_diff(
         else:
             raise NotImplementedError
 
-        loss_fn = loss_class(actor, value, loss_critic_type="l2")
+        loss_fn = loss_class(
+            actor,
+            value,
+            loss_critic_type="l2",
+            device=device,
+        )
 
         params = TensorDict.from_module(loss_fn, as_module=True)
 
@@ -9595,6 +9615,7 @@ def test_ppo_value_clipping(
                     value,
                     loss_critic_type="l2",
                     clip_value=clip_value,
+                    device=device,
                 )
 
         else:
@@ -9603,6 +9624,7 @@ def test_ppo_value_clipping(
                 value,
                 loss_critic_type="l2",
                 clip_value=clip_value,
+                device=device,
             )
             advantage(td)
             if composite_action_dist:

torchrl/objectives/ppo.py

@@ -440,7 +440,9 @@ def __init__(
                 raise ValueError(
                     f"clip_value must be a float or a scalar tensor, got {clip_value}."
                 )
-        self.register_buffer("clip_value", clip_value)
+            self.register_buffer("clip_value", clip_value.to(device))
+        else:
+            self.clip_value = None
         try:
             log_prob_keys = self.actor_network.log_prob_keys
             action_keys = self.actor_network.dist_sample_keys