[BUGFIX] Type fix for large tensors (#20922)

* Dimension type fix

* Review suggestion

Author: DominikaJedynak

src/operator/subgraph/dnnl/dnnl_conv.cc

@@ -59,11 +59,11 @@ static void UpdateConvWeightBias(NDArray* weight,
   const float* var_ptr     = variance.data().dptr<float>();
   DType* update_weight_ptr = update_weight.data().dptr<DType>();
   DType* update_bias_ptr   = update_bias.data().dptr<DType>();
-  size_t channel           = gamma.shape()[0];
+  index_t channel          = static_cast<index_t>(gamma.shape()[0]);
   const auto wshape        = weight->shape();
   size_t offset            = wshape.ProdShape(1, wshape.ndim());
 #pragma omp parallel for num_threads(engine::OpenMP::Get()->GetRecommendedOMPThreadCount())
-  for (int c = 0; c < static_cast<int>(channel); ++c) {
+  for (index_t c = 0; c < channel; ++c) {
     const DType* p1 = weight_ptr + c * offset;
     DType* p2       = update_weight_ptr + c * offset;
     float alpha     = (param->fix_gamma ? 1.0f : gamma_ptr[c]) / sqrt(var_ptr[c] + param->eps);

src/operator/subgraph/dnnl/dnnl_fc.cc

@@ -235,15 +235,15 @@ void SgDNNLFCOp::Forward(const OpContext& ctx,
     const mxnet::TShape oshape = output.shape();
     dnnl::memory::dims out_dims(2);
     if (oshape.ndim() == 2) {
-      out_dims[0] = static_cast<int>(oshape[0]);
-      out_dims[1] = static_cast<int>(oshape[1]);
+      out_dims[0] = static_cast<index_t>(oshape[0]);
+      out_dims[1] = static_cast<index_t>(oshape[1]);
     } else {
       if (!default_param.flatten) {
-        out_dims[0] = static_cast<int>(oshape.ProdShape(0, oshape.ndim() - 1));
-        out_dims[1] = static_cast<int>(oshape[oshape.ndim() - 1]);
+        out_dims[0] = static_cast<index_t>(oshape.ProdShape(0, oshape.ndim() - 1));
+        out_dims[1] = static_cast<index_t>(oshape[oshape.ndim() - 1]);
       } else {
-        out_dims[0] = static_cast<int>(static_cast<int>(oshape[0]));
-        out_dims[1] = static_cast<int>(oshape.ProdShape(1, oshape.ndim()));
+        out_dims[0] = static_cast<index_t>(oshape[0]);
+        out_dims[1] = static_cast<index_t>(oshape.ProdShape(1, oshape.ndim()));
       }
     }
     dnnl::memory::desc out_md =