Commit 1dc1b6fa authored by doubleHU's avatar doubleHU Committed by binmakeswell
Browse files

[NFC] polish...

[NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/cross_entropy_layer.h code style (#957)
parent 0e922da8
......@@ -8,8 +8,9 @@
#include "cuda_util.h"
template <typename T> class CrossEntropyLayer {
public:
template <typename T>
class CrossEntropyLayer {
public:
CrossEntropyLayer(float epsilon, int padding_idx, int max_batch_tokens);
virtual ~CrossEntropyLayer();
......@@ -22,7 +23,7 @@ public:
void set_cur_batch_shape(int batch_size, int seq_len, int vocab_size);
private:
private:
void allocate_mem_buffer() {
// allocate local gpu memory
_loss_buffer = cuda_malloc<float>(_max_batch_tokens * 2);
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment