github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-08 21:03:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/cross_entropy_layer.h code style (#957)

Browse Source
This commit is contained in:
doubleHU 2022-05-15 09:00:09 +08:00 committed by binmakeswell
parent 0e922da874
commit 1dc1b6fa00
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/kernel/cuda_native/csrc/kernels/include/cross_entropy_layer.h
View File

@ -8,8 +8,9 @@
#include "cuda_util.h" #include "cuda_util.h"
template <typename T> class CrossEntropyLayer { template <typename T>
public: class CrossEntropyLayer {
public:
CrossEntropyLayer(float epsilon, int padding_idx, int max_batch_tokens); CrossEntropyLayer(float epsilon, int padding_idx, int max_batch_tokens);
virtual ~CrossEntropyLayer(); virtual ~CrossEntropyLayer();
@ -22,7 +23,7 @@ public:
void set_cur_batch_shape(int batch_size, int seq_len, int vocab_size); void set_cur_batch_shape(int batch_size, int seq_len, int vocab_size);
private: private:
void allocate_mem_buffer() { void allocate_mem_buffer() {
// allocate local gpu memory // allocate local gpu memory
_loss_buffer = cuda_malloc<float>(_max_batch_tokens * 2); _loss_buffer = cuda_malloc<float>(_max_batch_tokens * 2);