ruff

Genentech · Aug 19, 2024 · 60f37cc · 60f37cc
1 parent 1a17828
commit 60f37cc
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/beignet/_segment_sum.py b/src/beignet/_segment_sum.py
@@ -48,8 +48,8 @@ def segment_sum(
     valid_indexes = indexes[valid_mask]
     valid_input = input[valid_mask]
 
-    output = torch.zeros(n, *input.shape[1:], device=input.device, dtype=valid_input.dtype)
-
-    return output.scatter_add(0, valid_indexes, valid_input).to(
-        **kwargs
+    output = torch.zeros(
+        n, *input.shape[1:], device=input.device, dtype=valid_input.dtype
     )
+
+    return output.scatter_add(0, valid_indexes, valid_input).to(**kwargs)