black applied

baler/modules/data_processing.py

@@ -57,9 +57,10 @@ def encoder_saver(model, model_path: str) -> None:
     Returns:
         None: Saved encoder state dictionary as `.pt` file.
     """
-    if hasattr(model.encoder,'state_dict'):
-        torch.save(model.encoder.state_dict(),model_path)
-    else: model.save_encoder(model_path)
+    if hasattr(model.encoder, "state_dict"):
+        torch.save(model.encoder.state_dict(), model_path)
+    else:
+        model.save_encoder(model_path)
 
 
 def decoder_saver(model, model_path: str) -> None:
@@ -72,9 +73,10 @@ def decoder_saver(model, model_path: str) -> None:
     Returns:
         None: Saved decoder state dictionary as `.pt` file.
     """
-    if hasattr(model.decoder,'state_dict'):
-        torch.save(model.decoder.state_dict(),model_path)
-    else: model.save_decoder(model_path)
+    if hasattr(model.decoder, "state_dict"):
+        torch.save(model.decoder.state_dict(), model_path)
+    else:
+        model.save_decoder(model_path)
 
 
 def initialise_model(model_name: str):

baler/modules/models.py

@@ -713,11 +713,11 @@ def get_final_layer_dims(self):
     def set_final_layer_dims(self, conv_op_shape):
         self.conv_op_shape = conv_op_shape
 
+
 class PJ_Conv_AE_FPGA(nn.Module):
     def __init__(self, n_features, z_dim=10, *args, **kwargs):
         super(PJ_Conv_AE_FPGA, self).__init__(*args, **kwargs)
 
-
         # Encoder layers
         self.en1 = nn.Conv2d(1, 20, kernel_size=5, stride=2, padding=2)
         self.en_act1 = nn.ReLU()
@@ -766,17 +766,27 @@ def forward(self, x):
         return decoded
 
     def get_final_layer_dims(self):
-        return 
+        return
 
     def set_final_layer_dims(self, conv_op_shape):
         self.conv_op_shape = conv_op_shape
 
     def save_encoder(self, file_path):
         # Create an instance of the encoder
-        encoder_instance = nn.Sequential(self.en1, self.en_act1, self.en2, self.en_act2, self.en3, self.en4)
+        encoder_instance = nn.Sequential(
+            self.en1, self.en_act1, self.en2, self.en_act2, self.en3, self.en4
+        )
         torch.save(encoder_instance.state_dict(), file_path)
 
     def save_decoder(self, file_path):
         # Create an instance of the decoder
-        decoder_instance = nn.Sequential(self.de1, self.de_act1, self.de2, self.de_unflatten, self.de_conv1, self.de_conv2, self.de_act2)
-        torch.save(decoder_instance.state_dict(), file_path)
+        decoder_instance = nn.Sequential(
+            self.de1,
+            self.de_act1,
+            self.de2,
+            self.de_unflatten,
+            self.de_conv1,
+            self.de_conv2,
+            self.de_act2,
+        )
+        torch.save(decoder_instance.state_dict(), file_path)