Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BugFix] TensorDictPrimer updates spec instead of overwriting #2332

Merged
merged 8 commits into from
Jul 31, 2024
Merged

[BugFix] TensorDictPrimer updates spec instead of overwriting #2332

Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
9a4718f
amend
matteobettini Jul 29, 2024
cfd30f8
amend
matteobettini Jul 29, 2024
4e3103d
amend
matteobettini Jul 29, 2024
abdaa61
amend
matteobettini Jul 30, 2024
54a0842
Merge remote-tracking branch 'origin/main' into primer
vmoens Jul 30, 2024
25e85e1
Merge remote-tracking branch 'origin/main' into primer
vmoens Jul 30, 2024
700f0be
test
matteobettini Jul 31, 2024
50940ee
Merge branch 'refs/heads/main' into primer
matteobettini Jul 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
amend
  • Loading branch information
@matteobettini
matteobettini committed Jul 29, 2024
commit 9a4718fc353bdbc41a34d9a35f9ca38fc895ea7d
16 changes: 7 additions & 9 deletions torchrl/envs/transforms/transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4638,15 +4638,13 @@ def transform_observation_spec(
raise ValueError(
f"observation_spec was expected to be of type CompositeSpec. Got {type(observation_spec)} instead."
)
for key, spec in self.primers.items():
if spec.shape[: len(observation_spec.shape)] != observation_spec.shape:
expanded_spec = self._expand_shape(spec)
spec = expanded_spec
try:
device = observation_spec.device
except RuntimeError:
device = self.device
observation_spec[key] = self.primers[key] = spec.to(device)
try:
device = observation_spec.device
matteobettini marked this conversation as resolved.
Show resolved Hide resolved
except RuntimeError:
device = self.device
if self.primers.shape[: len(observation_spec.shape)] != observation_spec.shape:
self.primers = self._expand_shape(self.primers)
observation_spec.update(self.primers.to(device))
return observation_spec

def transform_input_spec(self, input_spec: TensorSpec) -> TensorSpec:
Expand Down