DOC Adds code comment to nn.Linear

thomasjpfan · thomasjpfan · commit c01ba70ab116 · 2021-05-18T12:10:51.000-04:00
diff --git a/torch/nn/modules/linear.py b/torch/nn/modules/linear.py
@@ -86,6 +86,8 @@ def __init__(self, in_features: int, out_features: int, bias: bool = True,
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
+        # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
+        # uniform(-1/sqrt(in_features), 1/sqrt(in_features))
         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)
