dim size fix for index.numel() == 0

Author: rusty1s

torch_scatter/utils/gen.py

@@ -1,6 +1,12 @@
 from itertools import repeat
 
 
+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
+
+
 def gen(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     dim = range(src.dim())[dim]  # Get real dim value.
 
@@ -12,8 +18,8 @@ def gen(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
 
     # Generate output tensor if not given.
     if out is None:
-        dim_size = index.max().item() + 1 if dim_size is None else dim_size
         out_size = list(src.size())
+        dim_size = maybe_dim_size(index, dim_size)
         out_size[dim] = dim_size
         out = src.new_full(out_size, fill_value)