Add log_figure and log_histogram to TensorBoardLogger (#743)

Summary:
Pull Request resolved: #743

The SummaryWriter supports these just fine, so no reason not to include them here.

Reviewed By: galrotem

Differential Revision: D54958596

fbshipit-source-id: 6b54b015460e51b6672962848fb4f69b72393eaa

torchtnt/utils/loggers/tensorboard.py

@@ -167,6 +167,17 @@ def log_images(self: TensorBoardLogger, *args: Any, **kwargs: Any) -> None:
         if writer:
             writer.add_images(*args, **kwargs)
 
+    def log_figure(self: TensorBoardLogger, *args: Any, **kwargs: Any) -> None:
+        """Add matplotlib figure to TensorBoard.
+
+        Args:
+            *args (Any): Positional arguments passed to SummaryWriter.add_figure
+            **kwargs(Any): Keyword arguments passed to SummaryWriter.add_figure
+        """
+        writer = self._writer
+        if writer:
+            writer.add_figure(*args, **kwargs)
+
     def log_audio(self: TensorBoardLogger, *args: Any, **kwargs: Any) -> None:
         """Add audio data to TensorBoard.
 
@@ -202,6 +213,16 @@ def log_scalars(
                 walltime=walltime,
             )
 
+    def log_histogram(self: TensorBoardLogger, *args: Any, **kwargs: Any) -> None:
+        """Add histogram to TensorBoard.
+
+        Args:
+            *args (Any): Positional arguments passed to SummaryWriter.add_histogram
+            **kwargs(Any): Keyword arguments passed to SummaryWriter.add_histogram
+        """
+        if self._writer:
+            self._writer.add_histogram(*args, **kwargs)
+
     def flush(self: TensorBoardLogger) -> None:
         """Writes pending logs to disk."""