[Performance] Use TensorDict._new_unsafe in step

Vincent Moens · Vincent Moens · commit e5cba04df393 · 2025-04-22T16:56:49.000+01:00
ghstack-source-id: 8a117fb Pull Request resolved: #2905
diff --git a/torchrl/envs/gym_like.py b/torchrl/envs/gym_like.py
@@ -343,7 +343,7 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
                     for key, val in TensorDict(obs_dict, []).items(True, True)
                 )
         else:
-            tensordict_out = TensorDict(
+            tensordict_out = TensorDict._new_unsafe(
                 obs_dict,
                 batch_size=tensordict.batch_size,
             )

Original file line number	Diff line number	Diff line change
`@@ -343,7 +343,7 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:`
`343`	`343`	`for key, val in TensorDict(obs_dict, []).items(True, True)`
`344`	`344`	`)`
`345`	`345`	`else:`
`346`		`- tensordict_out = TensorDict(`
	`346`	`+ tensordict_out = TensorDict._new_unsafe(`
`347`	`347`	`obs_dict,`
`348`	`348`	`batch_size=tensordict.batch_size,`
`349`	`349`	`)`