[BugFix] Fix DDPG squeezing (#1487)

Signed-off-by: Matteo Bettini <matbet@meta.com>
pytorch · Sep 7, 2023 · d8504e7 · d8504e7
1 parent 786020d
commit d8504e7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/torchrl/objectives/ddpg.py b/torchrl/objectives/ddpg.py
@@ -11,10 +11,10 @@
 from typing import Tuple
 
 import torch
+
 from tensordict.nn import dispatch, make_functional, repopulate_module, TensorDictModule
 from tensordict.tensordict import TensorDict, TensorDictBase
 from tensordict.utils import NestedKey
-
 from torchrl.modules.tensordict_module.actors import ActorCriticWrapper
 from torchrl.objectives.common import LossModule
 from torchrl.objectives.utils import (
@@ -269,7 +269,6 @@ def forward(self, tensordict: TensorDictBase) -> TensorDict:
         """
         loss_value, td_error, pred_val, target_value = self._loss_value(tensordict)
         td_error = td_error.detach()
-        td_error = td_error.unsqueeze(tensordict.ndimension())
         if tensordict.device is not None:
             td_error = td_error.to(tensordict.device)
         tensordict.set(