Skip to content

Commit

Permalink
#9837: Assign workers after performing ref count cleanup in async mode
Browse files Browse the repository at this point in the history
  - This handles cases where a device tensor is reassigned to a host tensor
  - Exposed during model cache generation which uses the following pattern:
      device_tensor = device_tensor.cpu()
  • Loading branch information
tt-asaigal committed Jul 4, 2024
1 parent fec343c commit 112ea79
Showing 1 changed file with 2 additions and 4 deletions.
6 changes: 2 additions & 4 deletions tt_eager/tensor/tensor.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -188,10 +188,9 @@ struct Tensor {
Tensor &operator=(const Tensor &other) {
// Don't self-assign
if (this->tensor_attributes != other.tensor_attributes) {
// Assign workers before cleanup
this->workers = other.workers;
// Update ref count for curr tensor_attr and deallocate if needed
perform_cleanup_for_async_mode();
this->workers = other.workers;
this->tensor_id = other.tensor_id;
this->tensor_attributes = other.tensor_attributes;
this->deallocate_through_destructor = other.deallocate_through_destructor;
Expand All @@ -209,10 +208,9 @@ struct Tensor {
Tensor &operator=(Tensor &&other) {
// Don't self assign
if (this->tensor_attributes != other.tensor_attributes) {
// Assign workers before cleanup
this->workers = std::move(other.workers);
// Update ref count for curr tensor_attr and deallocate if needed
perform_cleanup_for_async_mode();
this->workers = std::move(other.workers);
this->tensor_id = std::move(other.tensor_id);
this->tensor_attributes = std::move(other.tensor_attributes);
this->deallocate_through_destructor = std::move(other.deallocate_through_destructor);
Expand Down

0 comments on commit 112ea79

Please sign in to comment.