BUGFIX: in JIT for CUDA backend

- Dimension check for linear index was incorrect