Extend DML max-length fix to other hardware (#296)

microsoft · Apr 23, 2024 · 54fe3fd · 54fe3fd
1 parent 2b9b666
commit 54fe3fd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/generators.cpp b/src/generators.cpp
@@ -68,8 +68,8 @@ std::unique_ptr<Search> CreateSearch(const GeneratorParams& params) {
 
 Generator::Generator(const Model& model, const GeneratorParams& params) : model_{model.shared_from_this()} {
 #if USE_DML
-  // Temporary fix to work around overflows for caches that are multiples of 4 on Intel hardware in DirectML
-  if (model.device_type_ == DeviceType::DML && model.IsIntelDevice() && params.search.max_length % 4 == 0) {
+  // Temporary fix to work around overflows for caches that are multiples of 4 in DirectML
+  if (model.device_type_ == DeviceType::DML && params.search.max_length % 4 == 0) {
     if (params.search.max_length == model.config_->model.context_length) {
       --const_cast<GeneratorParams&>(params).search.max_length;
     } else {