From 41955edc8dadb1456507cb6eaa8c6a5667b96e4b Mon Sep 17 00:00:00 2001 From: lizz Date: Sun, 12 Jan 2020 13:30:46 +0800 Subject: [PATCH] fix: dim shoud be int Signed-off-by: lizz --- torch_scatter/utils/gen.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torch_scatter/utils/gen.py b/torch_scatter/utils/gen.py index 45dee97b..aa6a6221 100644 --- a/torch_scatter/utils/gen.py +++ b/torch_scatter/utils/gen.py @@ -8,7 +8,8 @@ def maybe_dim_size(index, dim_size=None): if dim_size is not None: return dim_size - return index.max().item() + 1 if index.numel() > 0 else 0 + dim = index.max().item() + 1 if index.numel() > 0 else 0 + return int(dim) def broadcast(src, index, dim):