[JAX] [B] Fixed Batcher in DBiasCastTranspose Primitive (#843)

fixed batcher in dbias_cast_transpose primitive
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

transformer_engine/jax/cpp_extensions.py

@@ -4298,7 +4298,7 @@ class DBiasCastTransposePrimitive(BasePrimitive):
         _check_valid_batch_dims(batch_dims)
         assert DBiasCastTransposePrimitive.outer_primitive is not None
         dz, amax, scale, scale_inv = batched_args
-        dz_bdim, _, amax_bdim, _, _ = batch_dims
+        dz_bdim, amax_bdim, _, _ = batch_dims
 
         # Minus batch dim.
         transpose_axis_boundary = _normalize_axis_boundary(transpose_axis_boundary, dz.ndim - 1)