fix shadow

cuda/fastermoe/smart_schedule.cpp

@@ -44,10 +44,9 @@ void _reduce_grad(
         long expert_size) {
     auto smgr = getCudaStreamManager(t.device().index());
 
-    auto torch_stream = c10::cuda::getCurrentCUDAStream().stream();
     cudaEvent_t evt_stash;
     cudaEventCreate(&evt_stash);
-    cudaEventRecord(evt_stash, torch_stream);
+    cudaEventRecord(evt_stash, smgr->torchStream());
     FMOE_SWE(smgr->stream(0), evt_stash);
     cudaEventDestroy(evt_stash);
 

fmoe/fastermoe/schedule.py

@@ -37,7 +37,7 @@ class MoEForward(Function):
                 try:
                     # To skip torch autograd's version check.
                     with torch.autograd.graph.saved_tensors_hooks(nothing, nothing):
-                        y0 = expert_fn(x, torch.tensor([x.shape[0]], dtype=torch.int64))
+                        y0 = expert_fn(x, torch.tensor([x.shape[0]], dtype=torch.int64), expert_idx)
                 except Exception as e:
                     # Ignore the error and fall back for compatibility to older
                     # versions of PyTorch