[BugFix] GAE warning when gamma/lmbda are tensors (#2838)

Co-authored-by: Louis Faury <louis.faury@helsing.ai>
(cherry picked from commit d561115)

Author: 2 people

torchrl/objectives/value/advantages.py

@@ -1284,8 +1284,18 @@ def __init__(
             skip_existing=skip_existing,
             device=device,
         )
-        self.register_buffer("gamma", torch.tensor(gamma, device=self._device))
-        self.register_buffer("lmbda", torch.tensor(lmbda, device=self._device))
+        self.register_buffer(
+            "gamma",
+            gamma.to(self._device)
+            if isinstance(gamma, Tensor)
+            else torch.tensor(gamma, device=self._device),
+        )
+        self.register_buffer(
+            "lmbda",
+            lmbda.to(self._device)
+            if isinstance(lmbda, Tensor)
+            else torch.tensor(lmbda, device=self._device),
+        )
         self.average_gae = average_gae
         self.vectorized = vectorized
         self.time_dim = time_dim