Просмотр исходного кода

CUDA: fix scratch malloced on non-main device (#3220)

Johannes Gäßler 2 лет назад
Родитель
Сommit
578d8c8f5c
1 измененных файлов с 1 добавлено и 0 удалено
  1. 1 0
      ggml-cuda.cu

+ 1 - 0
ggml-cuda.cu

@@ -6970,6 +6970,7 @@ void ggml_cuda_assign_scratch_offset(struct ggml_tensor * tensor, size_t offset)
         return;
     }
     if (g_scratch_buffer == nullptr) {
+        ggml_cuda_set_device(g_main_device);
         CUDA_CHECK(cudaMalloc(&g_scratch_buffer, g_scratch_size));
     }