                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 rows = in->shape().elements() / in->shape().back();
+  size_t cols = in->shape().back();
+
   size_t colsToCopy = indices.size();
 
   int threads = std::min(MAX_THREADS, (int)colsToCopy);

                            Tensor beta,
                            float eps) {
  cudaSetDevice(adj->getDevice());
  int rows = y->shape()[0] * y->shape()[2] * y->shape()[3];
  int cols = y->shape()[1];

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

void CopyCols(Tensor out, const Tensor in, const std::vector<size_t>& indices) {
  cudaSetDevice(out->getDevice());

  size_t rows = in->shape().elements() / in->shape().back();
  size_t cols = in->shape().back();

  size_t colsToCopy = indices.size();

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