more logs

neuromodulation · Nov 12, 2024 · 2e77a48 · 2e77a48
1 parent 2326ab5
commit 2e77a48
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/speech_training_demo.py b/speech_training_demo.py
@@ -788,7 +788,9 @@ def prepare_dataset(batch):
                     "lr": torch.tensor(optimizer.param_groups[0]["lr"]),
                     "temp": torch.tensor(gumbel_temperature),
                     "grad_norm": torch.tensor(grad_norm),
-                    "Outputs": outputs
+                    "Outputs": outputs,
+                    "projectec_states": torch.tensor(outputs.projected_states),
+                    "projected_quant_states": torch.tensor(outputs.projected_quantized_states)
                 }
 
                 log_str = ""
@@ -851,7 +853,7 @@ def prepare_dataset(batch):
             val_logs["val_contrastive_loss"] += outputs.contrastive_loss
             val_logs["val_diversity_loss"] += outputs.diversity_loss
             val_logs["val_num_losses"] += batch["mask_time_indices"].sum()
-            val_logs["val_outputs"] += outputs
+            # val_logs["val_outputs"] += outputs
 
         # sum over devices in multi-processing
         if accelerator.num_processes > 1: