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: 0cd65696a91d83674212ca9a62dce02d1cabf44d
Pull Request resolved: #954
  • Loading branch information
vmoens committed Sep 16, 2024
1 parent 257c859 commit c8df23a
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 @@ -2323,18 +2323,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 c8df23a

Please sign in to comment.