From 6957bc38a112f10cd0bbef576de97719d9fce1a5 Mon Sep 17 00:00:00 2001
From: Yulu Jia <62443743+pleasantrabbit@users.noreply.github.com>
Date: Mon, 18 Jan 2021 01:02:53 -0800
Subject: [PATCH] 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>
---
 byteps/torch/ops.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/byteps/torch/ops.cc b/byteps/torch/ops.cc
index d3e589fe5..1a9864558 100644
--- 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