Custom metrics from compute_loss in TGA (2nd try)

parlai/core/torch_generator_agent.py

@@ -717,6 +717,12 @@ def _encoder_input(self, batch):
         """
         return self._model_input(batch)
 
+    def record_per_token_metrics(self, batch, loss_per_token):
+        """
+        Override this method for custom loss values that require loss_per_token.
+        """
+        pass
+
     def compute_loss(self, batch, return_output=False):
         """
         Compute and return the loss for the given batch.
@@ -752,6 +758,7 @@ def compute_loss(self, batch, return_output=False):
         self.record_local_metric(
             'token_em', AverageMetric.many(num_tokens_correct == num_target_tokens)
         )
+        self.record_per_token_metrics(batch, loss_per_token)
 
         # actually do backwards loss
         loss = loss_per_token.sum(dim=1)