浏览代码

fix backward

Denis Mazur 4 年之前
父节点
当前提交
97654aceb5
共有 1 个文件被更改,包括 1 次插入1 次删除
  1. 1 1
      hivemind/moe/client/expert.py

+ 1 - 1
hivemind/moe/client/expert.py

@@ -157,7 +157,7 @@ class _RemoteModuleCall(torch.autograd.Function):
         ]
 
         grad_inputs = cls.run_coroutine(
-            ctx.stub.rpc_backward(runtime_pb2.ExpertRequest(uid=ctx.uid, tensors=serialized_tensors)),
+            ctx.stub.rpc_backward(async_generate(runtime_pb2.ExpertRequest(uid=ctx.uid, tensors=serialized_tensors))),
         )
 
         deserialized_grad_inputs = [deserialize_torch_tensor(tensor) for tensor in grad_inputs.tensors]