(layer_id)
| 819 | ) |
| 820 | |
| 821 | def get_layer(layer_id): |
| 822 | return GLMBlock( |
| 823 | self.hidden_size, |
| 824 | self.num_attention_heads, |
| 825 | self.layernorm_epsilon, |
| 826 | layer_id, |
| 827 | inner_hidden_size=self.inner_hidden_size, |
| 828 | hidden_size_per_attention_head=self.hidden_size_per_attention_head, |
| 829 | layernorm=LayerNorm, |
| 830 | use_bias=True, |
| 831 | params_dtype=self.params_dtype, |
| 832 | position_encoding_2d=self.position_encoding_2d, |
| 833 | ) |
| 834 | |
| 835 | self.layers = torch.nn.ModuleList( |
| 836 | [get_layer(layer_id) for layer_id in range(self.num_layers)] |