diff --git a/pdebench/data_download/download_direct.py b/pdebench/data_download/download_direct.py
index 5cd0ee4..5a87249 100644
--- a/pdebench/data_download/download_direct.py
+++ b/pdebench/data_download/download_direct.py
@@ -50,9 +50,7 @@ def parse_metadata(pde_names):
         "3d_cfd",
     ]
 
-    assert all(
-            name.lower() in pde_list for name in pde_names
-    ), "PDE name not defined."
+    assert all(name.lower() in pde_list for name in pde_names), "PDE name not defined."
 
     # Filter the files to be downloaded
     meta_df["PDE"] = meta_df["PDE"].str.lower()
diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config_2D.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config_2D.yaml
index acd486e..8909062 100644
--- a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config_2D.yaml
+++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config_2D.yaml
@@ -1,4 +1,4 @@
-save: '../save/ReacDiff/'
+save: "../save/ReacDiff/"
 dt_save: 0.01
 ini_time: 0.
 fin_time: 1.
@@ -8,9 +8,9 @@ xL: 0.
 xR: 6.28318530718
 yL: 0.
 yR: 6.28318530718
-nu : 1.e0
+nu: 1.e0
 rho: 1.e0
 CFL: 2.5e-1
 if_show: 1
 show_steps: 100
-init_mode: 'react'
+init_mode: "react"
diff --git a/pdebench/data_gen/data_gen_NLE/config/config.yaml b/pdebench/data_gen/data_gen_NLE/config/config.yaml
index c79ac76..25f1620 100644
--- a/pdebench/data_gen/data_gen_NLE/config/config.yaml
+++ b/pdebench/data_gen/data_gen_NLE/config/config.yaml
@@ -9,7 +9,7 @@ hydra:
     dir: .
 
 args:
-  type: 'ReacDiff'  # 'advection'/'ReacDiff'/'burgers'/'CFD'
+  type: "ReacDiff"  # "advection"/"ReacDiff"/"burgers"/"CFD"
   dim: 1
   bd: "periodic"
   nbatch: 1000
diff --git a/pdebench/data_gen/velocity2vorticity.py b/pdebench/data_gen/velocity2vorticity.py
index 8181dc4..3081ced 100644
--- a/pdebench/data_gen/velocity2vorticity.py
+++ b/pdebench/data_gen/velocity2vorticity.py
@@ -92,8 +92,7 @@ def convert_velocity() -> None:
                 [vx[..., None], vy[..., None], vz[..., None]], axis=-1
             )
 
-            vorticity = compute_spectral_vorticity_jnp(
-                jnp.array(velocity), dx, dy, dz)
+            vorticity = compute_spectral_vorticity_jnp(jnp.array(velocity), dx, dy, dz)
 
             outfile["omega_x"][i] = np.array(vorticity[..., 0])
             outfile["omega_y"][i] = np.array(vorticity[..., 1])
diff --git a/pdebench/models/fno/train.py b/pdebench/models/fno/train.py
index 25a6912..2381c45 100644
--- a/pdebench/models/fno/train.py
+++ b/pdebench/models/fno/train.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 import torch
-from torch import nn
+import torch.nn as nn
 
 from pdebench.models.fno.fno import FNO1d, FNO2d, FNO3d
 from pdebench.models.fno.utils import FNODatasetMult, FNODatasetSingle
diff --git a/pdebench/models/fno/utils.py b/pdebench/models/fno/utils.py
index a98cbe9..4151a27 100644
--- a/pdebench/models/fno/utils.py
+++ b/pdebench/models/fno/utils.py
@@ -149,7 +149,6 @@
 from __future__ import annotations
 
 import math as mt
-import os
 from pathlib import Path
 
 import h5py
diff --git a/pdebench/models/unet/utils.py b/pdebench/models/unet/utils.py
index 4175b10..a39f564 100644
--- a/pdebench/models/unet/utils.py
+++ b/pdebench/models/unet/utils.py
@@ -149,7 +149,6 @@
 from __future__ import annotations
 
 import math as mt
-import os
 from pathlib import Path
 
 import h5py