   auto mp_indices = allocator->alloc<size_t>(indices.size());
   mp_indices->insert(indices.data(), indices.size());
 
+  int axisGPU = axis + ShapeGPU::size() - out->shape().size();
   gSelect<<<blocks, threads>>>(out->data(),
                                out->shape(),
                                in->data(),
                                in->shape(),
-                               axis,
+                               axisGPU,
                                mp_indices->data<size_t>());
 
   allocator->free(mp_indices);

               const std::vector<size_t>& indices) {
  cudaSetDevice(out->getDevice());

  size_t rows = out->shape()[0] * out->shape()[2] * out->shape()[3];
  size_t cols = out->shape()[1];
  size_t colsToCopy = indices.size();

  int threads = std::min(MAX_THREADS, (int)colsToCopy);

void LSTMOutputForward(Tensor out, std::vector<Tensor> inputs) {
  cudaSetDevice(out->getDevice());

  int rows = out->shape().elements() / out->shape().back();
  int cols = out->shape().back();

  int blocks = std::min(MAX_BLOCKS, rows);
  int threads = std::min(MAX_THREADS, cols);
