cleanup

Author: lucidrains

audiolm_pytorch/trainer.py

@@ -902,6 +902,7 @@ def data_tuple_to_kwargs(self, data):
             assert not has_duplicates(self.ds_fields), 'dataset fields must not have duplicate field names'
 
         return dict(zip(self.ds_fields, data))
+
     @contextmanager
     def wandb_tracker(self, project, run = None, hps = None):
         assert self.use_wandb_tracking, '`use_wandb_tracking` must be set to True on SemanticTransformerTrainer'
@@ -919,6 +920,7 @@ def wandb_tracker(self, project, run = None, hps = None):
         yield
 
         self.accelerator.end_training()
+
     def train_step(self):
         device = self.device
 
@@ -1197,6 +1199,7 @@ def wandb_tracker(self, project, run = None, hps = None):
         yield
 
         self.accelerator.end_training()  
+
     @property
     def device(self):
         return self.accelerator.device
@@ -1475,6 +1478,7 @@ def print(self, msg):
 
     def generate(self, *args, **kwargs):
         return self.train_wrapper.generate(*args, **kwargs)
+
     @contextmanager
     def wandb_tracker(self, project, run = None, hps = None):
         assert self.use_wandb_tracking, '`use_wandb_tracking` must be set to True on FineTransformerTrainer'
@@ -1491,7 +1495,8 @@ def wandb_tracker(self, project, run = None, hps = None):
 
         yield
 
-        self.accelerator.end_training()  
+        self.accelerator.end_training() 
+
     @property
     def device(self):
         return self.accelerator.device