diff --git a/src/models.py b/src/models.py
index b91078f..f814d6d 100644
--- a/src/models.py
+++ b/src/models.py
@@ -97,9 +97,9 @@ def __init__(self, partition_configs):
         self.fin1_act = nn.LeakyReLU()
         self.final_lin_transform = nn.Linear(50, 1)
         self.final_lin_act = nn.LeakyReLU()
-        #with torch.no_grad():
+        # with torch.no_grad():
         #    self.final_lin_transform.weight.copy_(torch.ones(self.n_partitions))
-        #for p in self.final_lin_transform.parameters():
+        # for p in self.final_lin_transform.parameters():
         #    p.requires_grad = False
 
     def init_conv_layers(self):
@@ -173,7 +173,8 @@ def forward(self, data_objects):
         output = self.final_lin_act(output)
 
         return {"output": output, "contributions": contributions}
-    
+
+
 class SlabGCN(nn.Module):
     """Class to customize the graph neural network."""
 
@@ -243,20 +244,20 @@ def __init__(self, partition_configs, global_config):
 
         # Hidden layers
         self.hidden_layers = nn.Sequential(
-                    *(
-                        [
-                            nn.Linear(self.hidden_size, self.hidden_size),
-                            nn.LeakyReLU(inplace=True),
-                            nn.Dropout(p=self.dropout),
-                        ]
-                        * (self.n_hidden - 1)
-                        + [
-                            nn.Linear(self.hidden_size, 1),
-                            nn.LeakyReLU(inplace=True),
-                            nn.Dropout(p=self.dropout),
-                        ]
-                    )
-                )
+            *(
+                [
+                    nn.Linear(self.hidden_size, self.hidden_size),
+                    nn.LeakyReLU(inplace=True),
+                    nn.Dropout(p=self.dropout),
+                ]
+                * (self.n_hidden - 1)
+                + [
+                    nn.Linear(self.hidden_size, 1),
+                    nn.LeakyReLU(inplace=True),
+                    nn.Dropout(p=self.dropout),
+                ]
+            )
+        )
 
     def init_conv_layers(self):
         """Initialize convolutional layers."""
@@ -320,7 +321,7 @@ def forward(self, data_objects):
         output = self.hidden_layers(hidden_data)
 
         return {"output": output}
-    
+
     def get_embeddings(self, data_objects):
         """Get the pooled embeddings of each partition.
 
@@ -358,6 +359,7 @@ def get_embeddings(self, data_objects):
 
         return embeddings
 
+
 if __name__ == "__main__":
     from pathlib import Path
 
diff --git a/workflows/basic_train_val_test.py b/workflows/basic_train_val_test.py
index a7857a0..3330d49 100644
--- a/workflows/basic_train_val_test.py
+++ b/workflows/basic_train_val_test.py
@@ -18,17 +18,19 @@
 dataset = AtomsDataset(root=dataset_path, prop_csv=prop_csv_path)
 
 # Process dataset
-dataset.process_data(layer_cutoffs=[3, 6],
-                         node_features=[
-                             ["atomic_number", "dband_center", "coordination"],
-                             ["atomic_number", "reactivity", "coordination"],
-                             ["atomic_number", "reactivity", "coordination"],
-                         ],
-                         edge_features=[
-                             ["bulk_bond_distance"],
-                             ["surface_bond_distance"],
-                             ["adsorbate_bond_distance"],
-                         ])
+dataset.process_data(
+    layer_cutoffs=[3, 6],
+    node_features=[
+        ["atomic_number", "dband_center", "coordination"],
+        ["atomic_number", "reactivity", "coordination"],
+        ["atomic_number", "reactivity", "coordination"],
+    ],
+    edge_features=[
+        ["bulk_bond_distance"],
+        ["surface_bond_distance"],
+        ["adsorbate_bond_distance"],
+    ],
+)
 
 # Create sampler
 sample_config = {"train": 0.8, "val": 0.1, "test": 0.1}