Clean up no_grad for trunc normal weight inits

timm/models/layers/weight_init.py

@@ -5,7 +5,7 @@ import warnings
 from torch.nn.init import _calculate_fan_in_and_fan_out
 
 
-def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+def _trunc_normal_(tensor, mean, std, a, b):
     # Cut & paste from PyTorch official master until it's in a few official releases - RW
     # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
     def norm_cdf(x):
@@ -17,7 +17,6 @@ def _no_grad_trunc_normal_(tensor, mean, std, a, b):
                       "The distribution of values may be incorrect.",
                       stacklevel=2)
 
-    with torch.no_grad():
     # Values are generated by using a truncated uniform distribution and
     # then using the inverse CDF for the normal distribution.
     # Get upper and lower cdf values
@@ -64,7 +63,8 @@ def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
         >>> w = torch.empty(3, 5)
         >>> nn.init.trunc_normal_(w)
     """
-    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+    with torch.no_grad():
+        return _trunc_normal_(tensor, mean, std, a, b)
 
 
 def trunc_normal_tf_(tensor, mean=0., std=1., a=-2., b=2.):
@@ -90,8 +90,8 @@ def trunc_normal_tf_(tensor, mean=0., std=1., a=-2., b=2.):
         >>> w = torch.empty(3, 5)
         >>> nn.init.trunc_normal_(w)
     """
-    _no_grad_trunc_normal_(tensor, 0, 1.0, a, b)
     with torch.no_grad():
+        _trunc_normal_(tensor, 0, 1.0, a, b)
         tensor.mul_(std).add_(mean)
     return tensor
 
@@ -111,9 +111,11 @@ def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
         # constant is stddev of standard normal truncated to (-2, 2)
         trunc_normal_tf_(tensor, std=math.sqrt(variance) / .87962566103423978)
     elif distribution == "normal":
+        with torch.no_grad():
             tensor.normal_(std=math.sqrt(variance))
     elif distribution == "uniform":
         bound = math.sqrt(3 * variance)
+        with torch.no_grad():
             tensor.uniform_(-bound, bound)
     else:
         raise ValueError(f"invalid distribution {distribution}")