code: def log_stats(
      self,
      epoch_number: int,
      batch_number: int,
      logs: Dict,
      layer_stats: Dict[str, Dict[str, float]],
  ) -> None:
      self.current_step += 1
      for layer, stats in layer_stats.items():
          for name, value in stats.items():
              self.writer.add_scalars(
                  f"gradient_stats/{layer}",
                  {name: value},
                  self.current_step,
              )
  
imports: import csv
  import os
  from typing import Dict, List, Optional, TextIO, Union
  import torch
  from poutyne.framework.callbacks._utils import atomic_lambda_save
  from poutyne.framework.callbacks.callbacks import Callback
label_desc: Use image node_id 2 for calling the TensorBoardGradientLogger obj's underlying member method code with example usage: obj.log_stats(epoch_number, batch_number, logs, layer_stats) without return types