Simple linear layer with weight initialization.
(rows, columns, init_method)
| 41 | |
| 42 | |
| 43 | def get_linear_layer(rows, columns, init_method): |
| 44 | """Simple linear layer with weight initialization.""" |
| 45 | layer = torch.nn.Linear(rows, columns) |
| 46 | init_method(layer.weight) |
| 47 | with torch.no_grad(): |
| 48 | layer.bias.zero_() |
| 49 | return layer |
| 50 | |
| 51 | @torch.jit.script |
| 52 | def gelu_impl(x): |