message: [Transpose] fix a bug that incorrectly computes the strides of the input/output tensors
diff: @@ -55,10 +55,15 @@ void Transpose::forward_kernel(TransposeMeta const *m,
  info.num_dim = out_domain.get_dim();
  assert(info.num_dim == m->num_dim);
  for (int i = 0; i < info.num_dim; i++) {
  - int in_dim_size = (in_domain.hi()[i] - in_domain.lo()[i] + 1);
  - int out_dim_size = (out_domain.hi()[i] - out_domain.lo()[i] + 1);
  - info.in_strides[i] = (i == 0) ? 1 : info.in_strides[i - 1] * in_dim_size;
  - info.out_strides[i] = (i == 0) ? 1 : info.out_strides[i - 1] * out_dim_size;
  + if (i == 0) {
  + info.in_strides[i] = 1;
  + info.out_strides[i] = 1;
  + } else {
  + int in_dim_size = (in_domain.hi()[i - 1] - in_domain.lo()[i - 1] + 1);
  + int out_dim_size = (out_domain.hi()[i - 1] - out_domain.lo()[i - 1] + 1);
  + info.in_strides[i] = info.in_strides[i - 1] * in_dim_size;
  + info.out_strides[i] = info.out_strides[i - 1] * out_dim_size;
  + }
  info.perm[i] = m->perm[i];
  }
  transpose_simple_kernel<<<GET_BLOCKS(out_domain.get_volume()),
  @@ -91,10 +96,15 @@ void Transpose::backward_kernel(TransposeMeta const *m,
  info.num_dim = in_grad_domain.get_dim();
  assert(info.num_dim == m->num_dim);
  for (int i = 0; i < info.num_dim; i++) {
  - int in_dim_size = (out_grad_domain.hi()[i] - out_grad_domain.lo()[i] + 1);
  - int out_dim_size = (in_grad_domain.hi()[i] - in_grad_domain.lo()[i] + 1);
  - info.in_strides[i] = (i == 0) ? 1 : info.in_strides[i - 1] * in_dim_size;
  - info.out_strides[i] = (i == 0) ? 1 : info.out_strides[i - 1] * out_dim_size;
  + if (i == 0) {
  + info.in_strides[i] = 1;
  + info.out_strides[i] = 1;
  + } else {
  + int in_dim_size = (out_grad_domain.hi()[i - 1] - out_grad_domain.lo()[i - 1] + 1);
  + int out_dim_size = (in_grad_domain.hi()[i - 1] - in_grad_domain.lo()[i - 1] + 1);
  + info.in_strides[i] = info.in_strides[i - 1] * in_dim_size;
  + info.out_strides[i] = info.out_strides[i - 1] * out_dim_size;
  + }
  info.perm[m->perm[i]] = i;
  }
  transpose_simple_kernel<<<GET_BLOCKS(in_grad_domain.get_volume()),
  