class_docstr: An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
  models.
class_docstr_tok: ["An","abstract","class","to","handle","weights","initialization","and","a","simple","interface","for","downloading","and","loading","pretrained","models","."]
docstr: Initialize the weights
docstr_tok: ["Initialize","the","weights"]
code: def _init_weights(
      self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]
  ) -> None:
      """Initialize the weights"""
      if isinstance(module, (nn.Linear, nn.Conv2d)):
          # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid
          # `trunc_normal_cpu` not implemented in `half` issues
          module.weight.data = nn.init.trunc_normal_(
              module.weight.data.to(torch.float32),
              mean=0.0,
              std=self.config.initializer_range,
          ).to(module.weight.dtype)
          if module.bias is not None:
              module.bias.data.zero_()
      elif isinstance(module, nn.LayerNorm):
          module.bias.data.zero_()
          module.weight.data.fill_(1.0)
  
imports: import math
  from typing import Dict, List, Optional, Set, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from ...activations import ACT2FN
  from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, SequenceClassifierOutput
  from ...modeling_utils import PreTrainedModel
  from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
  from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
  from .configuration_audio_spectrogram_transformer import ASTConfig
label_desc: Use image node_id 1 for calling the ASTPreTrainedModel obj's underlying member method code with example usage: obj._init_weights(module) without return types