fix

ztxz16 · May 31, 2024 · 3b76e7a · 3b76e7a
1 parent 1d373ba
commit 3b76e7a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/models/basellm.cpp b/src/models/basellm.cpp
@@ -66,6 +66,9 @@ namespace fastllm {
             this->kv.push_back(std::make_pair(Data(dataType), Data(dataType)));
         }
         for (int i = 0; i < kv->size(); i++) {
+            (*kv)[i].first.ToDevice(DataDevice::CPU);
+            (*kv)[i].second.ToDevice(DataDevice::CPU);
+
             this->kv[i].first.CopyFrom((*kv)[i].first);
             this->kv[i].second.CopyFrom((*kv)[i].second);
         }