From bec7a6bd8c8b44b4819433524e788e6a7c21433d Mon Sep 17 00:00:00 2001
From: Junjie Zhang <junzhang@nvidia.com>
Date: Fri, 15 Dec 2023 00:39:11 -0800
Subject: [PATCH] Fix softmax tensor declaration in ctor

---
 HugeCTR/src/layers/softmax_layer.cu | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/HugeCTR/src/layers/softmax_layer.cu b/HugeCTR/src/layers/softmax_layer.cu
index 662c1637ce..9e518fe732 100644
--- a/HugeCTR/src/layers/softmax_layer.cu
+++ b/HugeCTR/src/layers/softmax_layer.cu
@@ -22,6 +22,7 @@
 #include <linalg/binary_op.cuh>
 #include <linalg/reduce.cuh>
 #include <linalg/unary_op.cuh>
+#include <network_buffer_channels.hpp>
 #include <utils.hpp>
 namespace HugeCTR {
 
@@ -36,14 +37,15 @@ SoftmaxLayer<T>::SoftmaxLayer(const core23::Tensor& input_tensor,
   dims_ = input_tensor.shape().dims();
   hidden_size_ = input_tensor.shape().size(dims_ - 1);
   n_rows_ = len_ / hidden_size_;
-  workspace23_ =
-      core23::Tensor({(int64_t)n_rows_}, core23::DataType(core23::ToScalarType<T>::value));
-  identity23_ =
-      core23::Tensor({(int64_t)hidden_size_}, core23::DataType(core23::ToScalarType<T>::value));
-  softmax_out23_ =
-      core23::Tensor(input_tensor.shape(), core23::DataType(core23::ToScalarType<T>::value));
+  core23::BufferParams buf_p{.channel = GetBlobsBufferChannel()};
+  auto param = (input_tensor.my_params().buffer_params(buf_p));
+  workspace23_ = core23::Tensor(
+      param.shape({(int64_t)n_rows_}).data_type(core23::DataType(core23::ToScalarType<T>::value)));
+  identity23_ = core23::Tensor(param.shape({(int64_t)hidden_size_})
+                                   .data_type(core23::DataType(core23::ToScalarType<T>::value)));
+  softmax_out23_ = core23::Tensor(param.shape(input_tensor.shape())
+                                      .data_type(core23::DataType(core23::ToScalarType<T>::value)));
 }
-
 template <typename T>
 void SoftmaxLayer<T>::initialize() {
   CudaDeviceContext context(get_device_id());