support more types

test/test_segment.py

@@ -5,12 +5,11 @@ import torch
 from torch.autograd import gradcheck
 from torch_scatter import segment_coo, segment_csr
 
-from .utils import tensor
+from .utils import tensor, dtypes
 
 reductions = ['add', 'mean', 'min', 'max']
 grad_reductions = ['add', 'mean']
 
-dtypes = [torch.float]
 devices = [torch.device('cuda')]
 
 tests = [

test/utils.py

 import torch
-from torch.testing import get_all_dtypes
-
-dtypes = get_all_dtypes()
-
-dtypes.remove(torch.half)
-dtypes.remove(torch.short)  # PyTorch scatter does not work on short types.
-dtypes.remove(torch.bool)
-if hasattr(torch, 'bfloat16'):
-    dtypes.remove(torch.bfloat16)
 
+dtypes = [torch.float, torch.double, torch.int, torch.long]
 grad_dtypes = [torch.float, torch.double]
 
 devices = [torch.device('cpu')]

torch_scatter/helpers.py

@@ -4,12 +4,12 @@ import torch
 def min_value(dtype):
     try:
         return torch.finfo(dtype).min
-    except AttributeError:
-        return torch.info(dtype).min
+    except TypeError:
+        return torch.iinfo(dtype).min
 
 
 def max_value(dtype):
     try:
         return torch.finfo(dtype).max
-    except AttributeError:
-        return torch.info(dtype).max
+    except TypeError:
+        return torch.iinfo(dtype).max