class_docstr: Wrapper to track the statistics of the weights' gradient per layer and log them in TensorBoard per epoch.
  
  Args:
      writer (~torch.utils.tensorboard.writer.SummaryWriter): The TensorBoard writer.
      keep_bias (bool): Either or not to log the bias of the network.
  
  Example:
      Using TensorBoardGradientTracker::
  
          from torch.utils.tensorboard import SummaryWriter
          from poutyne import Model, TensorBoardGradientTracker
  
          writer = SummaryWriter('runs')
          tb_tracker = TensorBoardGradientTracker(writer)
  
          model = Model(...)
          model.fit_generator(..., callbacks=[tb_tracker])
class_docstr_tok: ["Wrapper","to","track","the","statistics","of","the","weights","'","gradient","per","layer","and","log","them","in","TensorBoard","per","epoch",".","Args",":","writer","(","~torch.utils.tensorboard.writer.SummaryWriter",")",":","The","TensorBoard","writer",".","keep_bias","(","bool",")",":","Either","or","not","to","log","the","bias","of","the","network",".","Example",":","Using","TensorBoardGradientTracker",":",":","from","torch.utils.tensorboard","import","SummaryWriter","from","poutyne","import","Model",",","TensorBoardGradientTracker","writer","=","SummaryWriter","(","'runs","'",")","tb_tracker","=","TensorBoardGradientTracker","(","writer",")","model","=","Model","(","...",")","model.fit_generator","(","...",",","callbacks=","[","tb_tracker","]",")"]
code: def _on_epoch_end_log(self, epoch_number: int, logs: Dict) -> None:
      gradient_distributions_stats = [
          "mean",
          "mean_std_dev_up",
          "mean_std_dev_down",
      ]
      other_gradient_stats = ["min", "max"]
  
      formatted_stats = self.tracker.get_stats(self.layer_names)
      for layer_name in self.layer_names:
          stats = formatted_stats[layer_name]
  
          for (
              gradient_distributions_stat
          ) in gradient_distributions_stats:
              self.writer.add_scalars(
                  f"gradient_distributions/{layer_name}",
                  {
                      gradient_distributions_stat: stats[
                          gradient_distributions_stat
                      ]
                  },
                  epoch_number,
              )
          for other_gradient_stat in other_gradient_stats:
              self.writer.add_scalars(
                  f"other_gradient_stats/{layer_name}",
                  {other_gradient_stat: stats[other_gradient_stat]},
                  epoch_number,
              )
  
imports: from typing import Dict, Iterable, List, Tuple
  import numpy
  import torch
  from poutyne.framework.callbacks.callbacks import Callback
label_desc: Use image node_id 2 for calling the TensorBoardGradientTracker obj's underlying member method code with example usage: obj._on_epoch_end_log(epoch_number, logs) without return types