Update LFNA version 1.0

2021-05-13 21:33:34 +08:00
parent 3d3a04705f
commit cfabd05de8
11 changed files with 340 additions and 299 deletions
--- a/lib/xlayers/weight_init.py
+++ b/lib/xlayers/weight_init.py
@@ -60,4 +60,7 @@ def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
      >>> w = torch.empty(3, 5)
      >>> nn.init.trunc_normal_(w)
    """
-    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+    if isinstance(tensor, list):
+        return [_no_grad_trunc_normal_(x, mean, std, a, b) for x in tensor]
+    else:
+        return _no_grad_trunc_normal_(tensor, mean, std, a, b)