Commit 5f463f9b authored by sangwzh's avatar sangwzh
Browse files

update devicetype check in RecordedCopy

parent 28286ab3
...@@ -223,7 +223,7 @@ void NDArray::RecordedCopyFromTo( ...@@ -223,7 +223,7 @@ void NDArray::RecordedCopyFromTo(
CHECK(from->ctx.device_type != to->ctx.device_type) CHECK(from->ctx.device_type != to->ctx.device_type)
<< "Recoding event is only called for the copy between CPU and GPU."; << "Recoding event is only called for the copy between CPU and GPU.";
CHECK(from->ctx.device_type == kDGLCUDA || to->ctx.device_type == kDGLROCM) CHECK(from->ctx.device_type == kDGLCUDA || to->ctx.device_type == kDGLROCM ||from->ctx.device_type == kDGLROCM||to->ctx.device_type == kDGLCUDA)
<< "At least one CUDA ctx needs to be involved."; << "At least one CUDA ctx needs to be involved.";
DeviceAPI::Get(kDGLCUDA)->RecordedCopyDataFromTo( DeviceAPI::Get(kDGLCUDA)->RecordedCopyDataFromTo(
......
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