diff --git a/colossalai/kernel/cuda_native/csrc/kernels/include/cuda_util.h b/colossalai/kernel/cuda_native/csrc/kernels/include/cuda_util.h index 1595257be..bc2258762 100644 --- a/colossalai/kernel/cuda_native/csrc/kernels/include/cuda_util.h +++ b/colossalai/kernel/cuda_native/csrc/kernels/include/cuda_util.h @@ -20,8 +20,7 @@ void check_gpu_error(T result, char const *const func, const char *const file, template void print_vec(const T *outv, std::string outn, int num_output_ele); -template -T *cuda_malloc(size_t ele_num); +template T *cuda_malloc(size_t ele_num); void cuda_free(void *pdata); @@ -29,6 +28,6 @@ template void check_nan_inf(const T *data_ptr, int dsize, bool check_nan_inf, std::string file, int line, cudaStream_t stream); -#define CHECK_NAN_INF(ptr, size, stream) \ - check_nan_inf((ptr), (size), true, __FILE__, __LINE__, (stream)); \ +#define CHECK_NAN_INF(ptr, size, stream) \ + check_nan_inf((ptr), (size), true, __FILE__, __LINE__, (stream)); \ check_nan_inf((ptr), (size), false, __FILE__, __LINE__, (stream)) diff --git a/colossalai/kernel/cuda_native/csrc/kernels/include/dropout.h b/colossalai/kernel/cuda_native/csrc/kernels/include/dropout.h index 336bbacc9..c2a4f7c20 100644 --- a/colossalai/kernel/cuda_native/csrc/kernels/include/dropout.h +++ b/colossalai/kernel/cuda_native/csrc/kernels/include/dropout.h @@ -1,15 +1,14 @@ #pragma once -#include #include #include #include +#include #include "kernels.h" -template -class Dropout { - public: +template class Dropout { +public: struct Config { float ratio; bool training; @@ -89,7 +88,7 @@ class Dropout { void SetTrainingMode(bool training) { _config.training = training; } - private: +private: uint8_t *_mask; Config _config; };