def reset_parameters(self) -> None:
"""Reset the parameters in FusedMLP module.
"""
for layer_idx in range(self.hidden_layer_num + 1):
flow.nn.init.kaiming_uniform_(self.weight(layer_idx),
a=math.sqrt(5))
(fan_in, _) = _calculate_fan_in_and_fan_out(self.weight(layer_idx))
bound = 1 / math.sqrt(fan_in)
flow.nn.init.uniform_(self.bias(layer_idx), -bound, bound)
def reset_parameters(self) -> None:
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
(fan_in, _) = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)