Use insert_or_assign instead of erase+emplace (#164868)

insert_or_assign does effectively the same thing as
erase+emplace but more efficiently since the search
does not need to be repeated

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164868
Approved by: https://github.com/eqy
This commit is contained in:
Lakshay Garg
2025-10-08 19:13:49 +00:00
committed by PyTorch MergeBot
parent 37c6087334
commit a4110fedcf
4 changed files with 4 additions and 8 deletions

View File

@ -138,8 +138,7 @@ struct TORCH_API ExecutionTraceObserver { // NOLINT
// So we need to remove the key and insert the key with the new value.
data_ptr_to_storage_id.erase(raw_data_ptr);
data_ptr_to_storage_id[raw_data_ptr] = id;
data_ptr_to_weak_storage_ptr.erase(raw_data_ptr);
data_ptr_to_weak_storage_ptr.emplace(
data_ptr_to_weak_storage_ptr.insert_or_assign(
raw_data_ptr, t_storage.getWeakStorageImpl());
return id;
} else {