Tensor: Fix use_ahwb_ flag and tests on local device involved.

PiperOrigin-RevId: 498249332
2022-12-28 14:29:43 -08:00
parent 9580f04571
commit 1924f1cdff
3 changed files with 22 additions and 36 deletions
@@ -458,7 +458,8 @@ void Tensor::TrackAhwbUsage(uint64_t source_location_hash) const {
      ahwb_tracking_key_ = tensor_internal::FnvHash64(ahwb_tracking_key_, dim);
    }
  }
-  use_ahwb_ = ahwb_usage_track_.contains(ahwb_tracking_key_);
+  // Keep flag value if it was set previously.
+  use_ahwb_ = use_ahwb_ || ahwb_usage_track_.contains(ahwb_tracking_key_);
 }

 #else  // MEDIAPIPE_TENSOR_USE_AHWB