How to make cuda functioning correctly in 3d grid?

Question

I coded a small cuda function as following:

e=8
@cuda.jit
def my_kernel(e):
    i,j = cuda.grid(2)

    if i <= e and j<=e:
        i**2+j**2

# Host code   

threadsperblock = 16,16
blockspergrid = 32,32
my_kernel[blockspergrid, threadsperblock](e)
print('ok')

Now if I turn it into 3d form it gives me an error:

e=8
@cuda.jit
def my_kernel(e):
    i,j,k = cuda.grid(3)

    if i <= e and j<=e and k<=e:
        i**2+j**2+k**2

# Host code   

threadsperblock = 16,16,16
blockspergrid = 32,32,32
my_kernel[blockspergrid, threadsperblock](e)
print('ok')

error:
[1] Call to cuLaunchKernel results in CUDA_ERROR_INVALID_VALUE

How to solve this? Thanks really much.