class_docstr: An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
  models.
code: def _init_weights(self, module):
      if isinstance(module, (nn.Linear, nn.Embedding)):
          module.weight.data.normal_(
              mean=0.0, std=self.config.initializer_range
          )
          if isinstance(module, nn.Linear) and module.bias is not None:
              module.bias.data.zero_()
      elif isinstance(module, nn.LayerNorm):
          module.bias.data.zero_()
          module.weight.data.fill_(1.0)
      elif isinstance(module, EinLinear):
          for i in range(module.n_models):
              nn.init.kaiming_uniform_(
                  module.weight[i],
                  a=math.sqrt(5)
                  / self.config.kaiming_initializer_range,
              )
              if module.bias is not None:
                  fan_in, _ = nn.init._calculate_fan_in_and_fan_out(
                      module.weight[i]
                  )
                  bound = (
                      1 / math.sqrt(fan_in)
                  ) * self.config.initializer_range
                  nn.init.uniform_(module.bias[i], -bound, bound)
  
imports: import math
  import os
  from dataclasses import dataclass
  from typing import Optional, Tuple
  import numpy
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import functional
  from ...modeling_utils import PreTrainedModel
  from ...utils import ModelOutput, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_trajectory_transformer import TrajectoryTransformerConfig
label_desc: Use image node_id 2 for calling the TrajectoryTransformerPreTrainedModel obj's underlying member method code with example usage: obj._init_weights(module) without return types