 
 void CudaDevice::Take(const Array& a, const Array& indices, int8_t axis, const Array& out) {
     CheckDevicesCompatible(a, indices, out);
+    CheckCudaError(cudaSetDevice(index()));
     VisitDtype(out.dtype(), [&](auto pt) {
         using T = typename decltype(pt)::type;
 

    if (bytesize == 0) {
        return;
    }
    if (&src_device == this || nullptr != dynamic_cast<CudaDevice*>(&src_device)) {
        // Copy between CUDA devices
        CheckCudaError(cudaMemcpy(dst, src, bytesize, cudaMemcpyDeviceToDevice));

    assert(out.ndim() == 1);
    assert(out.shape()[0] > 0);

    CheckCudaError(cudaSetDevice(index()));
    VisitDtype(out.dtype(), [&](auto pt) {
        using T = typename decltype(pt)::type;
        int64_t n = out.shape()[0];
