code: def get_weight_drop_module(name: str, weight_dropout, **kwargs):
      return {"Conv2d": WeightDropConv2d, "Linear": WeightDropLinear}[
          name
      ](weight_dropout, **kwargs)
  
imports: import copy
  import warnings
  from typing import List, Optional, cast, Dict
  from baal.bayesian.common import replace_layers_in_module, _patching_wrapper, BayesianModule
  import torch
  from torch import nn
  from packaging.version import parse
label_desc: Use image node_id 1 for calling a global function with example usage: get_weight_drop_module(name, weight_dropout) and returns: dict