From c2e63870945822b3a305b0ffc4bd486cc1b77a82 Mon Sep 17 00:00:00 2001 From: Stefan Oehmcke Date: Sat, 27 Jul 2024 23:52:57 +0200 Subject: [PATCH] Fix SparseGlobalAvgPool in pool.py Fixing issue of SparseGlobalAvgPool giving back only the first feature. For clearer code I also changed max to amax so the magic indexing vanishes. --- spconv/pytorch/pool.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/spconv/pytorch/pool.py b/spconv/pytorch/pool.py index 23845074..becf42d9 100644 --- a/spconv/pytorch/pool.py +++ b/spconv/pytorch/pool.py @@ -270,9 +270,9 @@ def forward(self, input: spconv.SparseConvTensor): real_inds = out_indices[i, :counts_cpu_np[i]] real_features = input.features[real_inds] if self.is_mean: - real_features_reduced = torch.mean(real_features, dim=0)[0] + real_features_reduced = torch.mean(real_features, dim=0) else: - real_features_reduced = torch.max(real_features, dim=0)[0] + real_features_reduced = torch.amax(real_features, dim=0) res_features_list.append(real_features_reduced) res = torch.stack(res_features_list) return res @@ -581,4 +581,4 @@ def __init__(self, ALL_POOL_LAYERS = set([ SparseAvgPool3d, SparseAvgPool2d, SparseAvgPool1d, SparseMaxPool1d, SparseMaxPool2d, SparseMaxPool3d, SparseMaxPool4d, SparseAvgPool, SparseMaxPool -]) \ No newline at end of file +])