Memory-efficient stats gathering (#227)

pytorch · Jun 26, 2022 · aad5d04 · aad5d04
1 parent f7177cc
commit aad5d04
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchrl/trainers/helpers/envs.py b/torchrl/trainers/helpers/envs.py
@@ -336,7 +336,8 @@ def get_stats_random_rollout(
     while n < cfg.init_env_steps:
         _td_stats = proof_environment.rollout(max_steps=cfg.init_env_steps)
         n += _td_stats.numel()
-        td_stats.append(_td_stats)
+        td_stats.append(_td_stats.to_tensordict().select(key).cpu())
+        del _td_stats
     td_stats = torch.cat(td_stats, 0)
 
     if key is None: