copy to avoid bug

microsoft · Dec 9, 2024 · 7a9b6bb · 7a9b6bb
1 parent 4b12d57
commit 7a9b6bb
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/generators.cpp b/src/generators.cpp
@@ -297,6 +297,7 @@ void Generator::ComputeLogits(DeviceSpan<int32_t> next_tokens) {
   if (computed_logits_)
     throw std::runtime_error("ComputeLogits called again without calling AppendTokens or GenerateNextToken first");
 
+  next_tokens.CopyDeviceToCpu();
   auto logits = state_->Run(search_->GetSequenceLength(), next_tokens, search_->GetNextIndices());
   if (g_log.enabled && g_log.model_logits) {
     auto& stream = Log("model_logits");