MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / __init__

Method __init__

tensorrt_llm/models/bert/model.py:66–92  ·  view source on GitHub ↗
(self,
                 hidden_size,
                 num_attention_heads,
                 max_position_embeddings,
                 hidden_act='relu',
                 tp_group=None,
                 tp_size=1,
                 dtype=None)

Source from the content-addressed store, hash-verified

64class BertEncoderLayer(Module):
65
66 def __init__(self,
67 hidden_size,
68 num_attention_heads,
69 max_position_embeddings,
70 hidden_act='relu',
71 tp_group=None,
72 tp_size=1,
73 dtype=None):
74 super().__init__()
75 self.input_layernorm = LayerNorm(normalized_shape=hidden_size,
76 dtype=dtype)
77
78 self.attention = BertAttention(
79 hidden_size=hidden_size,
80 num_attention_heads=num_attention_heads,
81 max_position_embeddings=max_position_embeddings,
82 tp_group=tp_group,
83 tp_size=tp_size,
84 dtype=dtype)
85 self.mlp = MLP(hidden_size=hidden_size,
86 ffn_hidden_size=hidden_size * 4,
87 hidden_act=hidden_act,
88 tp_group=tp_group,
89 tp_size=tp_size,
90 dtype=dtype)
91 self.post_layernorm = LayerNorm(normalized_shape=hidden_size,
92 dtype=dtype)
93
94 def forward(self,
95 hidden_states,

Callers

nothing calls this directly

Calls 4

LayerNormClass · 0.50
BertAttentionClass · 0.50
MLPClass · 0.50
__init__Method · 0.45

Tested by

no test coverage detected