'Gets weight decay for index.
Returns 0 for non-weights if the name of weights are provided for `__init__`.
Parameters
index : int
The index for weight.
Returns
wd : float
Weight decay for this index.'

def _get_wd(self, index):

    wd = self.wd
    if (index in self.param_dict):
        wd *= self.param_dict[index].wd_mult
    elif (index in self.wd_mult):
        wd *= self.wd_mult[index]
    elif (index in self.idx2name):
        wd *= self.wd_mult.get(self.idx2name[index], 1.0)
    return wd
