From 6331eb3e2f5d113515ad4b1693160c7677b3b991 Mon Sep 17 00:00:00 2001
From: LXXXXR <73265258+LXXXXR@users.noreply.github.com>
Date: Fri, 18 Dec 2020 17:00:47 +0800
Subject: [PATCH] [Bug] Fix bug in collect_results (#114)

* fix bug in collect results

* fix compatibility issue in build.yml

* fix compatibility issue in build.yml

* revert changes in ci

* add comments
---
 mmcls/apis/test.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/mmcls/apis/test.py b/mmcls/apis/test.py
index dd7738274..b8d1c3472 100644
--- a/mmcls/apis/test.py
+++ b/mmcls/apis/test.py
@@ -106,7 +106,11 @@ def collect_results_cpu(result_part, size, tmpdir=None):
         part_list = []
         for i in range(world_size):
             part_file = osp.join(tmpdir, f'part_{i}.pkl')
-            part_list.append(mmcv.load(part_file))
+            part_result = mmcv.load(part_file)
+            # When data is severely insufficient, an empty part_result
+            # on a certain gpu could makes the overall outputs empty.
+            if part_result:
+                part_list.append(part_result)
         # sort the results
         ordered_results = []
         for res in zip(*part_list):
@@ -140,8 +144,11 @@ def collect_results_gpu(result_part, size):
     if rank == 0:
         part_list = []
         for recv, shape in zip(part_recv_list, shape_list):
-            part_list.append(
-                pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
+            part_result = pickle.loads(recv[:shape[0]].cpu().numpy().tobytes())
+            # When data is severely insufficient, an empty part_result
+            # on a certain gpu could makes the overall outputs empty.
+            if part_result:
+                part_list.append(part_result)
         # sort the results
         ordered_results = []
         for res in zip(*part_list):