CUDA: fix scratch malloced on non-main device

This commit is contained in:
JohannesGaessler 2023-09-16 21:48:29 +02:00
parent b541b4f0b1
commit 391dab70d7

View file

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