code: def train(self, mode=True):
          r
          super().train(mode)
  
          if mode:
              mu.apply_leaf(self, mu.set_train_mode)
  
          return self
docs: r"""
          Sets the module in training mode.
  
          This has any effect only on certain modules. See documentations of
          particular modules for details of their behaviors in training/evaluation
          mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
          etc.
  
          Returns:
              Module: self