diff --git a/setup.py b/setup.py
index d893d39..a40202e 100644
--- a/setup.py
+++ b/setup.py
@@ -1,32 +1,30 @@
 from setuptools import setup, find_packages
 
 setup(
-    name='xlb',
-    version='0.2.1',
-    description='XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML',
-    long_description=open('README.md').read(),
-    long_description_content_type='text/markdown',
-    author='Mehdi Ataei',
-    url='https://github.com/Autodesk/XLB',
-    license='Apache License 2.0',
+    name="xlb",
+    version="0.2.1",
+    description="XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    author="Mehdi Ataei",
+    url="https://github.com/Autodesk/XLB",
+    license="Apache License 2.0",
     packages=find_packages(),
     install_requires=[
-        'matplotlib>=3.9.2',
-        'numpy>=2.1.2',
-        'pyvista>=0.44.1',
-        'trimesh>=4.4.9',
-        'warp-lang>=1.4.0',
-        'numpy-stl>=3.1.2',
-        'pydantic>=2.9.1',
-        'ruff>=0.6.5',
-        'jax>=0.4.34'  # Base JAX CPU-only requirement
+        "matplotlib>=3.9.2",
+        "numpy>=2.1.2",
+        "pyvista>=0.44.1",
+        "trimesh>=4.4.9",
+        "warp-lang>=1.4.0",
+        "numpy-stl>=3.1.2",
+        "pydantic>=2.9.1",
+        "ruff>=0.6.5",
+        "jax>=0.4.34",  # Base JAX CPU-only requirement
     ],
     extras_require={
-        'cuda': ['jax[cuda12]>=0.4.34'],  # For CUDA installations
-        'tpu': ['jax[tpu]>=0.4.34'],  # For TPU installations
+        "cuda": ["jax[cuda12]>=0.4.34"],  # For CUDA installations
+        "tpu": ["jax[tpu]>=0.4.34"],  # For TPU installations
     },
-    python_requires='>=3.10',
-    dependency_links=[
-        'https://storage.googleapis.com/jax-releases/libtpu_releases.html'
-    ],
+    python_requires=">=3.10",
+    dependency_links=["https://storage.googleapis.com/jax-releases/libtpu_releases.html"],
 )
diff --git a/xlb/operator/boundary_condition/bc_do_nothing.py b/xlb/operator/boundary_condition/bc_do_nothing.py
index 56a332f..0ff6775 100644
--- a/xlb/operator/boundary_condition/bc_do_nothing.py
+++ b/xlb/operator/boundary_condition/bc_do_nothing.py
@@ -76,4 +76,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
\ No newline at end of file
+        return f_post
diff --git a/xlb/operator/boundary_condition/bc_equilibrium.py b/xlb/operator/boundary_condition/bc_equilibrium.py
index 77f408f..260ceb5 100644
--- a/xlb/operator/boundary_condition/bc_equilibrium.py
+++ b/xlb/operator/boundary_condition/bc_equilibrium.py
@@ -101,4 +101,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
\ No newline at end of file
+        return f_post
diff --git a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
index 38657e5..f47fe2f 100644
--- a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
+++ b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
@@ -205,4 +205,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
\ No newline at end of file
+        return f_post
diff --git a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
index 19a3013..bac9387 100644
--- a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
+++ b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
@@ -86,4 +86,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
\ No newline at end of file
+        return f_post
diff --git a/xlb/operator/boundary_condition/bc_grads_approximation.py b/xlb/operator/boundary_condition/bc_grads_approximation.py
index 94ddba3..217c65c 100644
--- a/xlb/operator/boundary_condition/bc_grads_approximation.py
+++ b/xlb/operator/boundary_condition/bc_grads_approximation.py
@@ -321,4 +321,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
\ No newline at end of file
+        return f_post