torch: fix hang after int tensor push_pull (#358)

mark task done after averaging an int tensor. This fixes a bug introduced in 46944e8. Signed-off-by: Yulu Jia <yulu.jia@bytedance.com>
bytedance · Jan 18, 2021 · 6957bc3 · 6957bc3
1 parent 46944e8
commit 6957bc3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/byteps/torch/ops.cc b/byteps/torch/ops.cc
@@ -81,6 +81,7 @@ void StartTask(::torch::Tensor tensor, ::torch::Tensor output, int average,
 #if TORCH_VERSION >= 1005000000
           if (isIntegralType(output.scalar_type())) {
             output.floor_divide_(byteps_size());
+            handle_manager.MarkDone(handle, status);
             return;
           }
 #endif