text: import torch
  import torch.nn as nn
  
  
  def init_weight(m):
      if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
          nn.init.kaiming_normal_(m.weight)
          if m.bias is not None:
              m.bias.data.zero_()
      elif isinstance(m, nn.BatchNorm2d):
          m.weight.data.fill_(1)
          m.bias.data.zero_()
      elif isinstance(m, nn.Linear):
          m.weight.data.normal_(0, 0.01)
          m.bias.data.zero_()