linting

torch_scatter/gather.py

@@ -5,8 +5,13 @@ from torch_scatter import segment_cpu, gather_cpu
 if torch.cuda.is_available():
     from torch_scatter import gather_cuda, segment_cuda
 
-gat = lambda is_cuda: gather_cuda if is_cuda else gather_cpu  # noqa
-seg = lambda is_cuda: segment_cuda if is_cuda else segment_cpu  # noqa
+
+def gat(is_cuda):
+    return gather_cuda if is_cuda else gather_cpu
+
+
+def seg(is_cuda):
+    return segment_cuda if is_cuda else segment_cpu
 
 
 class GatherCOO(torch.autograd.Function):

torch_scatter/segment.py

@@ -6,8 +6,13 @@ from torch_scatter.helpers import min_value, max_value
 if torch.cuda.is_available():
     from torch_scatter import segment_cuda, gather_cuda
 
-seg = lambda is_cuda: segment_cuda if is_cuda else segment_cpu  # noqa
-gat = lambda is_cuda: gather_cuda if is_cuda else gather_cpu  # noqa
+
+def seg(is_cuda):
+    return segment_cuda if is_cuda else segment_cpu
+
+
+def gat(is_cuda):
+    return gather_cuda if is_cuda else gather_cpu
 
 
 class SegmentCOO(torch.autograd.Function):