Skip to content

Commit

Permalink
[Performance] Remove list against list check during set
Browse files Browse the repository at this point in the history
ghstack-source-id: 9fb608e476b16a323ebbe8198e2ca94463007f58
Pull Request resolved: #954
  • Loading branch information
vmoens committed Aug 9, 2024
1 parent dd80ecc commit c4636a1
Showing 1 changed file with 10 additions and 12 deletions.
22 changes: 10 additions & 12 deletions tensordict/_td.py
Original file line number Diff line number Diff line change
Expand Up @@ -2254,18 +2254,16 @@ def _set_at_str(self, key, value, idx, *, validated, non_blocking: bool):

if is_non_tensor(value) and not (self._is_shared or self._is_memmap):
dest = tensor_in
is_diff = dest[idx].tolist() != value.tolist()
if is_diff:
dest_val = dest.maybe_to_stack()
dest_val[idx] = value
if dest_val is not dest:
self._set_str(
key,
dest_val,
validated=True,
inplace=False,
ignore_lock=True,
)
dest_val = dest.maybe_to_stack()
dest_val[idx] = value
if dest_val is not dest:
self._set_str(
key,
dest_val,
validated=True,
inplace=False,
ignore_lock=True,
)
return

if isinstance(idx, tuple) and len(idx) and isinstance(idx[0], tuple):
Expand Down

0 comments on commit c4636a1

Please sign in to comment.