修复多分类问题 (#20)

zhihu · May 29, 2019 · 235542d · 235542d
1 parent 7253884
commit 235542d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/cuBERT/op_out/AdditionalOutputLayer.cpp b/src/cuBERT/op_out/AdditionalOutputLayer.cpp
@@ -13,9 +13,11 @@ namespace cuBERT {
         this->handle = handle;
         this->hidden_size = hidden_size;
         this->num_labels = num_labels;
+
+        std::cerr << "hidden_size:" << this->hidden_size << " num_labels:" << this->num_labels << std::endl;
 
-        this->output_weights = static_cast<T *>(cuBERT::malloc(sizeof(T) * hidden_size));
-        cuBERT::memcpy(this->output_weights, output_weights, sizeof(T) * hidden_size, 1);
+        this->output_weights = static_cast<T *>(cuBERT::malloc(sizeof(T) * hidden_size * this->num_labels));
+        cuBERT::memcpy(this->output_weights, output_weights, sizeof(T) * hidden_size * this->num_labels, 1);
 
         if (output_bias != nullptr) {
             this->output_bias = static_cast<T *>(cuBERT::malloc(sizeof(T) * num_labels * max_batch_size));
@@ -25,6 +27,7 @@ namespace cuBERT {
         } else {
             this->output_bias = nullptr;
         }
+
     }
 
     template <typename T>

diff --git a/src/cuBERT/op_out/AdditionalOutputLayer.h b/src/cuBERT/op_out/AdditionalOutputLayer.h
@@ -4,6 +4,9 @@
 
 #include <cstddef>
 
+
+#include "cuBERT/op/Softmax.h"
+
 namespace cuBERT {
 
     /**
-Original file line number
+Diff line change
@@ Expand Up / @@ -4,6 +4,9 @@ @@
     #include <cstddef>
+    #include "cuBERT/op/Softmax.h"
     namespace cuBERT {
         /**
@@ Expand Down @@