(normalized_shape, eps=1e-5, elementwise_affine=True, export=False)
| 68 | |
| 69 | |
| 70 | def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False): |
| 71 | if not export and torch.cuda.is_available(): |
| 72 | try: |
| 73 | from apex.normalization import FusedLayerNorm |
| 74 | return FusedLayerNorm(normalized_shape, eps, elementwise_affine) |
| 75 | except ImportError: |
| 76 | pass |
| 77 | return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine) |
| 78 | |
| 79 | |
| 80 | def Linear(in_features, out_features, bias=True): |