[Lint] Black linter

endure/lcm/model/kaplsm_model.py

@@ -7,11 +7,13 @@
 
 DECISION_DIM = 64
 
+
 class KapEmbedding(nn.Module):
     """
     Special embedding that creates separate embeddings for each K_i on each
     level. Number of k's will dictate the number of linear layers.
     """
+
     def __init__(
         self,
         input_size: int,
@@ -37,6 +39,7 @@ def forward(self, x: Tensor) -> Tensor:
 
         return out
 
+
 class KapModel(nn.Module):
     def __init__(
         self,
@@ -90,7 +93,8 @@ def _split_input(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
             capacities = capacities.to(torch.long)
             capacities = F.one_hot(capacities, num_classes=self.capacity_range)
         else:
-            capacities = torch.unflatten(capacities, 1, (-1, self.capacity_range))
+            capacities = torch.unflatten(
+                capacities, 1, (-1, self.capacity_range))
 
         size_ratio = capacities[:, 0, :]
         k_cap = capacities[:, 1:, :]

endure/lsm/solver/klsm_solver.py

@@ -8,6 +8,7 @@
 from .util import kl_div_con, get_bounds
 from .util import H_DEFAULT, T_DEFAULT, LAMBDA_DEFAULT, ETA_DEFAULT, K_DEFAULT
 
+
 class KLSMSolver:
     def __init__(self, config: dict[str, Any]):
         self.config = config
@@ -28,8 +29,10 @@ def robust_objective(
         lamb, eta = x[-2:]
         design = LSMDesign(h=h, T=t, K=kaps, policy=Policy.KHybrid)
         query_cost = 0
-        query_cost += z0 * kl_div_con((self.cf.Z0(design, system) - eta) / lamb)
-        query_cost += z1 * kl_div_con((self.cf.Z1(design, system) - eta) / lamb)
+        query_cost += z0 * \
+            kl_div_con((self.cf.Z0(design, system) - eta) / lamb)
+        query_cost += z1 * \
+            kl_div_con((self.cf.Z1(design, system) - eta) / lamb)
         query_cost += q * kl_div_con((self.cf.Q(design, system) - eta) / lamb)
         query_cost += w * kl_div_con((self.cf.W(design, system) - eta) / lamb)
         cost = eta + (rho * lamb) + (lamb * query_cost)