"docs/vscode:/vscode.git/clone" did not exist on "f911287cc94218b3d65d97c233ed34d5b729c8c5"
Unverified Commit 5265631d authored by Vladimir's avatar Vladimir Committed by GitHub
Browse files

use a correct device when creating OptionalCUDAGuard (#2583)

parent 2832e7b9
......@@ -34,7 +34,7 @@ void swap_blocks(
char *dst_ptr = static_cast<char*>(dst.data_ptr());
const int64_t block_size_in_bytes = src.element_size() * src[0].numel();
const at::cuda::OptionalCUDAGuard device_guard(src_device);
const at::cuda::OptionalCUDAGuard device_guard(src_device.is_cuda() ? src_device : dst_device);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
// NOTE(woosuk): This can be slow if the number of blocks is large.
for (const auto& pair : block_mapping) {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment