[Test] Fix warnings in tests

ghstack-source-id: d4ed75d
Pull Request resolved: #2886

Author: Vincent Moens

test/test_cost.py

@@ -181,6 +181,9 @@
     pytest.mark.filterwarnings(
         "ignore:dep_util is Deprecated. Use functions from setuptools instead"
     ),
+    pytest.mark.filterwarnings(
+        "ignore:The PyTorch API of nested tensors is in prototype"
+    ),
 ]
 
 
@@ -16679,9 +16682,13 @@ def test_hf(self, from_text):
         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
         tokenizer.pad_token = tokenizer.eos_token
 
-        model = OPTForCausalLM(OPTConfig())
+        model = OPTForCausalLM(OPTConfig()).eval()
         policy_inference = TransformersWrapper(
-            model, tokenizer=tokenizer, generate=True, from_text=from_text
+            model,
+            tokenizer=tokenizer,
+            generate=True,
+            from_text=from_text,
+            return_log_probs=True,
         )
         policy_train = TransformersWrapper(
             model, tokenizer=tokenizer, generate=False, from_text=False

torchrl/envs/custom/chess.py

@@ -606,8 +606,8 @@ def _step(self, tensordict):
         reward = torch.tensor([reward_val], dtype=torch.float32)
         dest.set("reward", reward)
         dest.set("turn", turn)
-        dest.set("done", [done])
-        dest.set("terminated", [done])
+        dest.set("done", torch.tensor([done]))
+        dest.set("terminated", torch.tensor([done]))
         if self.pixels:
             dest.set("pixels", self._get_tensor_image(board=self.board))
         return dest

torchrl/objectives/ppo.py

@@ -584,7 +584,10 @@ def _log_weight(
             self.tensor_keys.sample_log_prob,
             adv_shape,
         )
-
+        if prev_log_prob is None:
+            raise KeyError(
+                f"Couldn't find the log-prob {self.tensor_keys.sample_log_prob} in the input data."
+            )
         if prev_log_prob.requires_grad:
             raise RuntimeError(
                 f"tensordict stored {self.tensor_keys.sample_log_prob} requires grad."