code: def on_train_begin(self, logs: Dict):
      self.layers = [
          n
          for n, p in self.model.network.named_parameters()
          if self._keep_layer(p, n)
      ]
  
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 GradientLoggerBase obj's underlying member method code with example usage: obj.on_train_begin(logs) without return types