Fix the segfault in the nvfp4 quantization (#2214)

* Fix the segfault in the nvfp4 quantization
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

transformer_engine/common/util/nvfp4_transpose.cuh

@@ -1433,7 +1433,8 @@ void nvfp4_quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *o
   const size_t block_size = THREADS_NUM;
 
   const size_t scale_stride = output->scale_inv.shape[1];
-  const size_t scale_stride_transpose = output->columnwise_scale_inv.shape[1];
+  const size_t scale_stride_transpose =
+      return_transpose ? output->columnwise_scale_inv.shape[1] : 0;
 
   nvfp4_scale_t *const scales_ptr = reinterpret_cast<nvfp4_scale_t *>(output->scale_inv.dptr);
   nvfp4_scale_t *const scales_transpose_ptr =