also now saves loss ( hahah funny meme)

| ||

|| |_

1 files changed, 2 insertions, 4 deletions

diff --git a/swr2_asr/train.py b/swr2_asr/train.py
index 6eaf4c1..9a8620f 100644
--- a/swr2_asr/train.py
+++ b/swr2_asr/train.py
@@ -351,10 +351,8 @@ def train(
                     [{batch_idx * len(spectrograms)}/{data_len} \
                     ({100.0 * batch_idx / len(train_loader)}%)]\t \
                     Loss: {loss.item()}"
-                    
             )
-        return loss  
-
+            return loss.item()
 
 def test(model, device, test_loader, criterion):
     """Test"""
@@ -476,7 +474,7 @@ def run(learning_rate: float = 5e-4, batch_size: int = 8, epochs: int = 3) -> No
             torch.save({
                 'epoch': epoch,
                 'model_state_dict': model.state_dict(),
-                 },MODEL_SAVE_PATH)
+                'loss': loss},MODEL_SAVE_PATH)
         test(model=model, device=device, test_loader=test_loader, criterion=criterion)