Skip to content

GitLab

Commit 881b10c2 authored by Rick Ho's avatar Rick Ho
Browse files

make it run

parent 691e92e1
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
pytorch/cuda/moe_cuda_kernel.cu
View file @ 881b10c2
...@@ -20,9 +20,9 @@ ...@@ -20,9 +20,9 @@
#define CEIL(_x_,_y_) (((_x_)-1)/(_y_)+1) #define CEIL(_x_,_y_) (((_x_)-1)/(_y_)+1)
#define MOE_DEBUG // #define MOE_DEBUG
// #define MOE_BREAKDOWN // #define MOE_BREAKDOWN
#define MOE_DEBUG_SCATTER // #define MOE_DEBUG_SCATTER
template <typename scalar_t> template <typename scalar_t>
__global__ __global__
...@@ -114,7 +114,7 @@ void moe_cuda_forward_impl( ...@@ -114,7 +114,7 @@ void moe_cuda_forward_impl(
for (int i = 0; i < batch_size; ++i) { for (int i = 0; i < batch_size; ++i) {
pos[i] = expert_ptr[gate[i]]++; pos[i] = expert_ptr[gate[i]]++;
} }
for (int i = batch_size - 1; i > 0; --i) { for (int i = tot_expert - 1; i > 0; --i) {
expert_ptr[i] = expert_ptr[i - 1]; expert_ptr[i] = expert_ptr[i - 1];
} }
expert_ptr[0] = 0; expert_ptr[0] = 0;
......
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