fix crash with return alternative on CUDA (#1733)

OpenNMT · Jun 27, 2024 · 451c27b · 451c27b
1 parent 72a461a
commit 451c27b
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/decoding.cc b/src/decoding.cc
@@ -70,9 +70,14 @@ namespace ctranslate2 {
                                                   StorageView& ids) {
     if (!decoder.output_layer_is_updated())
       return;
+    ctranslate2::Device device = ids.device();
+    if (device != Device::CPU)
+      ids = ids.to(Device::CPU);
     auto* ids_data = ids.data<int32_t>();
     for (dim_t i = 0; i < ids.size(); ++i)
       ids_data[i] = decoder.to_original_word_id(ids_data[i]);
+    if (ids.device() != device)
+      ids = ids.to(device);
   }
 
   template <typename T>