Set headscale false (#4364) (#4366)

Signed-off-by: MaximumEntropy <sandeep.subramanian.1@umontreal.ca> Signed-off-by: smajumdar <smajumdar@nvidia.com>
NVIDIA · Jun 17, 2022 · 106a097 · 106a097
1 parent 3ebe9ae
commit 106a097
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 3 deletions.
diff --git a/nemo/collections/asr/models/classification_models.py b/nemo/collections/asr/models/classification_models.py
@@ -481,6 +481,7 @@ def training_step(self, batch, batch_nb):
 
         self.log('train_loss', loss_value)
         self.log('learning_rate', self._optimizer.param_groups[0]['lr'])
+        self.log('global_step', self.trainer.global_step)
 
         self._accuracy(logits=logits, labels=labels)
         topk_scores = self._accuracy.compute()

diff --git a/nemo/collections/asr/models/ctc_models.py b/nemo/collections/asr/models/ctc_models.py
@@ -559,7 +559,11 @@ def training_step(self, batch, batch_nb):
             log_probs=log_probs, targets=transcript, input_lengths=encoded_len, target_lengths=transcript_len
         )
 
-        tensorboard_logs = {'train_loss': loss_value, 'learning_rate': self._optimizer.param_groups[0]['lr']}
+        tensorboard_logs = {
+            'train_loss': loss_value,
+            'learning_rate': self._optimizer.param_groups[0]['lr'],
+            'global_step': self.trainer.global_step,
+        }
 
         if hasattr(self, '_trainer') and self._trainer is not None:
             log_every_n_steps = self._trainer.log_every_n_steps

diff --git a/nemo/collections/asr/models/label_models.py b/nemo/collections/asr/models/label_models.py
@@ -277,6 +277,7 @@ def training_step(self, batch, batch_idx):
 
         self.log('loss', loss)
         self.log('learning_rate', self._optimizer.param_groups[0]['lr'])
+        self.log('global_step', self.trainer.global_step)
 
         self._accuracy(logits=logits, labels=labels)
         top_k = self._accuracy.compute()

diff --git a/nemo/collections/asr/models/rnnt_models.py b/nemo/collections/asr/models/rnnt_models.py
@@ -687,7 +687,11 @@ def training_step(self, batch, batch_nb):
                 log_probs=joint, targets=transcript, input_lengths=encoded_len, target_lengths=target_length
             )
 
-            tensorboard_logs = {'train_loss': loss_value, 'learning_rate': self._optimizer.param_groups[0]['lr']}
+            tensorboard_logs = {
+                'train_loss': loss_value,
+                'learning_rate': self._optimizer.param_groups[0]['lr'],
+                'global_step': self.trainer.global_step,
+            }
 
             if (sample_id + 1) % log_every_n_steps == 0:
                 self.wer.update(encoded, encoded_len, transcript, transcript_len)

diff --git a/nemo/collections/asr/models/ssl_models.py b/nemo/collections/asr/models/ssl_models.py
@@ -469,7 +469,10 @@ def training_step(self, batch, batch_nb):
             spectrograms, spec_masks, encoded, encoded_len, targets, target_lengths
         )
 
-        tensorboard_logs = {'learning_rate': self._optimizer.param_groups[0]['lr']}
+        tensorboard_logs = {
+            'learning_rate': self._optimizer.param_groups[0]['lr'],
+            'global_step': self.trainer.global_step,
+        }
 
         for loss_name, loss_val in loss_val_dict.items():
             tensorboard_logs['train_' + loss_name] = loss_val