ruff

PriorLabs · Feb 4, 2025 · dbfc9ea · dbfc9ea
1 parent 3a5e033
commit dbfc9ea
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/tabpfn/model/encoders.py b/src/tabpfn/model/encoders.py
@@ -120,7 +120,11 @@ def select_features(x: torch.Tensor, sel: torch.Tensor) -> torch.Tensor:
     if B == 1:
         return x[:, :, sel[0]]
 
-    new_x = torch.zeros((sequence_length, B, total_features), device=x.device, dtype=x.dtype)
+    new_x = torch.zeros(
+        (sequence_length, B, total_features),
+        device=x.device,
+        dtype=x.dtype,
+    )
 
     # For each batch, compute the number of selected features.
     sel_counts = sel.sum(dim=-1)  # shape: (B,)