Fix flake8 errors

Author: denix56

benchmarks/test_benchmark_nearest.py

@@ -31,6 +31,7 @@
 NEAREST_GROUPS = [1, 2, 4, 8, 16, 32]
 FEATURES = [8, 64, 200]
 
+
 def _assert_nearest_within_cuda(
     out_cuda: torch.Tensor,
     out_triton: torch.Tensor,

torch_cluster/triton/_kernels.py

@@ -3,7 +3,6 @@
 
 import triton
 import triton.language as tl
-from triton.language.extra import libdevice
 
 
 def _prune_knn_configs(configs, args, **kwargs):
@@ -452,7 +451,11 @@ def _knn_segmented_kernel(
                     padding_option="zero",
                 )
             if COSINE:
-                prod = tl.dot(x, y, input_precision=INPUT_PRECISION)  # MxV dot.
+                prod = tl.dot(
+                    x,
+                    y,
+                    input_precision=INPUT_PRECISION,
+                )  # MxV dot.
                 acc_dot += tl.sum(prod, axis=1)
                 acc_x_sq += tl.sum(x * x, axis=1)
                 y_block_ptr = tl.advance(y_block_ptr, (BLOCK_D, 0))
@@ -489,7 +492,11 @@ def _knn_segmented_kernel(
             can_reorder=True
         )  # Merge buffers.
         sorted_key = tl.sort(combo_key, descending=False)  # Sort 2K keys.
-        best_dist_key = tl.gather(sorted_key, k_offsets, axis=0)  # Keep best K.
+        best_dist_key = tl.gather(
+            sorted_key,
+            k_offsets,
+            axis=0,
+        )  # Keep best K.
 
     k_mask = k_offsets < K  # Mask for real k.
     best_dist_key = tl.where(k_mask, best_dist_key, INF_KEY)

torch_cluster/triton/knn.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 from typing import Optional
 
-import torch
 from torch import Tensor
 
 from .segmented import segmented_topk_search

torch_cluster/triton/segmented.py

@@ -2,7 +2,6 @@
 
 from typing import Optional
 
-import triton
 import torch
 from torch import Tensor