DOC Adds code comment to clarify nn.Linear.reset_parameters (#58487)

thomasjpfan · facebook-github-bot · commit 145a6f798517 · 2021-05-20T06:15:47.000-07:00
Summary: Fixes #57109 Adds comment to clarify `a=sqrt(5)` in `nn.Linear.reset_parameters`. Pull Request resolved: #58487 Reviewed By: ailzhang Differential Revision: D28548391 Pulled By: jbschlosser fbshipit-source-id: 2d5910b2576a04f19edbd8b8515cdb55fc249ce5
diff --git a/torch/nn/modules/linear.py b/torch/nn/modules/linear.py
@@ -86,6 +86,9 @@ 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)). For details, see
+        # https://github.com/pytorch/pytorch/issues/57109
         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)
