amend

pytorch · Jun 21, 2024 · f1b6aad · f1b6aad
1 parent 8b2c2aa
commit f1b6aad
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/torchrl/modules/distributions/continuous.py b/torchrl/modules/distributions/continuous.py
@@ -2,6 +2,8 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
 import warnings
 from numbers import Number
 from typing import Dict, Optional, Sequence, Tuple, Union

diff --git a/torchrl/objectives/sac.py b/torchrl/objectives/sac.py
@@ -647,7 +647,7 @@ def _qvalue_v1_loss(
         self, tensordict: TensorDictBase
     ) -> Tuple[Tensor, Dict[str, Tensor]]:
         target_params = self._cached_target_params_actor_value
-        with set_exploration_type(self.deterministic_mode):
+        with set_exploration_type(self.deterministic_sampling_mode):
             target_value = self.value_estimator.value_estimate(
                 tensordict, target_params=target_params
             ).squeeze(-1)