From 7bc601d71589109bd0e7a244ff34dafa232f400e Mon Sep 17 00:00:00 2001
From: xmy <2081932959@qq.com>
Date: Wed, 15 Jan 2025 13:22:22 +0800
Subject: [PATCH] Fix GPU detection in PerStoreFeatureNode(#17592)

---
 src/meta_schedule/feature_extractor/per_store_feature.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/meta_schedule/feature_extractor/per_store_feature.cc b/src/meta_schedule/feature_extractor/per_store_feature.cc
index 82bc7c2de078..955e39df4fc3 100644
--- a/src/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/meta_schedule/feature_extractor/per_store_feature.cc
@@ -1392,7 +1392,8 @@ class PerStoreFeatureNode : public FeatureExtractorNode {
 
   Array<runtime::NDArray> ExtractFrom(const TuneContext& tune_context,
                                       const Array<MeasureCandidate>& candidates) {
-    bool is_gpu = tune_context->target.value()->kind->name == "cuda";
+    auto& target_keys = tune_context->target.value()->keys;
+    bool is_gpu = std::find(target_keys.begin(), target_keys.end(), "gpu") != target_keys.end();
     std::vector<runtime::NDArray> results;
     results.resize(candidates.size());
     std::unique_ptr<tir::group6::Feature> feature_group6 = nullptr;