[BugFix] Fix KL import errors (#1207)

pytorch · May 30, 2023 · 7ed63a2 · 7ed63a2
1 parent e8f5efe
commit 7ed63a2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrl/envs/transforms/rlhf.py b/torchrl/envs/transforms/rlhf.py
@@ -13,8 +13,8 @@
     repopulate_module,
 )
 from tensordict.utils import _normalize_key, is_seq_of_nested_key
-from torchrl.data import CompositeSpec, UnboundedContinuousTensorSpec
-from torchrl.envs import Transform
+from torchrl.data.tensor_specs import CompositeSpec, UnboundedContinuousTensorSpec
+from torchrl.envs.transforms.transforms import Transform
 
 
 class KLRewardTransform(Transform):