[Docs] Fix warnings in mkdocs build (continued) (#24740)
Signed-off-by: Zerohertz <ohg3417@gmail.com>
This commit is contained in:
@@ -143,16 +143,8 @@ class AriaProjector(nn.Module):
|
||||
projects ViT's outputs into MoE's inputs.
|
||||
|
||||
Args:
|
||||
patch_to_query_dict (dict): Maps patch numbers to their corresponding
|
||||
query numbers,
|
||||
e.g., {1225: 128, 4900: 256}. This allows for different query sizes
|
||||
based on image resolution.
|
||||
embed_dim (int): Embedding dimension.
|
||||
num_heads (int): Number of attention heads.
|
||||
kv_dim (int): Dimension of key and value.
|
||||
ff_dim (int): Hidden dimension of the feed-forward network.
|
||||
output_dim (int): Output dimension.
|
||||
norm_layer (nn.Module): Normalization layer. Default is nn.LayerNorm.
|
||||
config: [AriaConfig](https://huggingface.co/docs/transformers/main/model_doc/aria#transformers.AriaConfig)
|
||||
containing projector configuration parameters.
|
||||
|
||||
Outputs:
|
||||
A tensor with the shape of (batch_size, query_number, output_dim)
|
||||
@@ -282,8 +274,8 @@ class AriaTextMoELayer(nn.Module):
|
||||
Forward pass of the MoE Layer.
|
||||
|
||||
Args:
|
||||
hidden_states (torch.Tensor): Input tensor of shape (batch_size,
|
||||
sequence_length, hidden_size).
|
||||
hidden_states: Input tensor of shape
|
||||
(batch_size, sequence_length, hidden_size).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor after passing through the MoE layer.
|
||||
|
||||
Reference in New Issue
Block a user