diff --git a/scripts/sbml_demo.py b/scripts/sbml_demo.py
index 3f14d51..d8b8de0 100644
--- a/scripts/sbml_demo.py
+++ b/scripts/sbml_demo.py
@@ -11,19 +11,22 @@
 reactions_sympy = sbml.sbml_to_sympy(model_sbml)
 sym_module = sbml.sympy_to_enzax(reactions_sympy)
 
-parameters_all = [({p.getId(): p.getValue() for p in r.getKineticLaw().getListOfParameters()}) for r in model_sbml.getListOfReactions()]
+parameters_all = [
+    ({p.getId(): p.getValue() for p in r.getKineticLaw().getListOfParameters()})
+    for r in model_sbml.getListOfReactions()
+]
 parameters = {}
 for i in parameters_all:
     parameters.update(i)
 
-compartments ={c.getId(): c.volume for c in model_sbml.getListOfCompartments()}
+compartments = {c.getId(): c.volume for c in model_sbml.getListOfCompartments()}
 
 species = [s.getId() for s in model_sbml.getListOfSpecies()]
 
 balanced_species_dict = {}
 unbalanced_species_dict = {}
 for i in model_sbml.getListOfSpecies():
-    if i.boundary_condition == False:
+    if not i.boundary_condition:
         balanced_species_dict.update({i.getId(): i.getInitialConcentration()})
     else:
         unbalanced_species_dict.update({i.getId(): i.getInitialConcentration()})
@@ -33,24 +36,38 @@
 
 para = {**parameters, **compartments, **unbalanced_species_dict}
 
-stoichmatrix  = jnp.zeros((model_sbml.getNumSpecies(), model_sbml.getNumReactions()), dtype=jnp.float64)
-i = 0 
+stoichmatrix = jnp.zeros(
+    (model_sbml.getNumSpecies(), model_sbml.getNumReactions()),
+    dtype=jnp.float64,
+)
+i = 0
 for reaction in model_sbml.getListOfReactions():
     for r in reaction.getListOfReactants():
-        stoichmatrix = stoichmatrix.at[species.index(r.getSpecies()), i].set(-int(r.getStoichiometry()))
+        stoichmatrix = stoichmatrix.at[species.index(r.getSpecies()), i].set(
+            -int(r.getStoichiometry())
+        )
     for p in reaction.getListOfProducts():
-        stoichmatrix = stoichmatrix.at[species.index(p.getSpecies()), i].set(int(p.getStoichiometry()))
+        stoichmatrix = stoichmatrix.at[species.index(p.getSpecies()), i].set(
+            int(p.getStoichiometry())
+        )
     i+=1
 
-structure = KineticModelStructure(stoichmatrix, jnp.array(balanced_ix), jnp.array(unbalanced_ix))
+structure = KineticModelStructure(
+    stoichmatrix, jnp.array(balanced_ix), jnp.array(unbalanced_ix)
+)
 
-kinmodel_sbml = KineticModelSbml(parameters=para, balanced_ids=balanced_species_dict, structure=structure, sym_module=sym_module)
+kinmodel_sbml = KineticModelSbml(
+    parameters=para,
+    balanced_ids=balanced_species_dict,
+    structure=structure,
+    sym_module=sym_module
+)
 
-y0 = jnp.array([2.,4])
+y0 = jnp.array([2.0, 4])
 kinmodel_sbml.flux(y0)
 
 kinmodel_sbml.dcdt(t=1, conc=y0)
 
-guess = jnp.full((2),0.01)
+guess = jnp.full((2), 0.01)
 steady_state = get_kinetic_model_steady_state(kinmodel_sbml, guess)
-print(steady_state)
\ No newline at end of file
+print(steady_state)
diff --git a/src/enzax/kinetic_model.py b/src/enzax/kinetic_model.py
index d6d3f3e..50b60d0 100644
--- a/src/enzax/kinetic_model.py
+++ b/src/enzax/kinetic_model.py
@@ -82,5 +82,9 @@ def flux(
         self,
         conc_balanced: Float[Array, " n_balanced"],
     ) -> Float[Array, " n"]:
-        flux = jnp.array(self.sym_module(**self.parameters, **dict(zip(self.balanced_ids, conc_balanced))))
-        return flux
\ No newline at end of file
+        flux = jnp.array(
+            self.sym_module(
+                **self.parameters, **dict(zip(self.balanced_ids, conc_balanced))
+            )
+        )
+        return flux