update

vitamin.py

@@ -31,7 +31,7 @@ from torch.utils.checkpoint import checkpoint
 from timm.models.layers import create_attn, get_norm_layer, get_norm_act_layer, create_conv2d, make_divisible, trunc_normal_tf_
 
 
-from timm.layers import to_2tuple, DropPath, Format, trunc_normal_
+from timm.layers import to_2tuple, DropPath, Format # , trunc_normal_
 from timm.layers.norm_act import _create_act
 from timm.models._registry import register_model
 from timm.models._manipulate import named_apply, checkpoint_seq
@@ -335,6 +335,39 @@ class HybridEmbed(nn.Module):
         x = x.flatten(2).transpose(1, 2)
         return x
 
+def _trunc_normal_(tensor, mean, std, a, b):
+    # rewrite timm trunc normal
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    l = norm_cdf((a - mean) / std)
+    u = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [l, u], then translate to
+    # [2l-1, 2u-1].
+    tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated standard normal
+    # tensor.erfinv_() # NOTE: deleted as "erfinv_cuda" not implemented for 'BFloat16'
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    tensor.clamp_(min=a, max=b)
+    return tensor
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    with torch.no_grad():
+        return _trunc_normal_(tensor, mean, std, a, b)
+
 class ViTamin(nn.Module):
     """ hack timm VisionTransformer
     """