From 54fe3fdff709ef39a46f37ff287fb062cf5c7332 Mon Sep 17 00:00:00 2001
From: Patrice Vignola <vignola.patrice@gmail.com>
Date: Tue, 23 Apr 2024 09:58:12 -0700
Subject: [PATCH] Extend DML max-length fix to other hardware (#296)

---
 src/generators.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/generators.cpp b/src/generators.cpp
index d7035b602..49751c3e9 100644
--- 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 {