utransformer

This commit is contained in:
Senstella
2025-09-28 19:22:54 +09:00
parent b9cc48bd25
commit 2761171fe3
2 changed files with 209 additions and 4 deletions

View File

@@ -214,7 +214,7 @@ class DINOv3ViTDropPath(nn.Module):
class DINOv3ViTMLP(nn.Module):
def __init__(self, config):
def __init__(self, config: DINOv3ViTConfig):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
@@ -236,7 +236,7 @@ class DINOv3ViTMLP(nn.Module):
class DINOv3ViTGatedMLP(nn.Module):
def __init__(self, config):
def __init__(self, config: DINOv3ViTConfig):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
@@ -274,8 +274,8 @@ class DINOv3ViTLayer(nn.Module):
self.mlp = DINOv3ViTMLP(config)
self.layer_scale2 = DINOv3ViTLayerScale(config)
def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None) -> torch.Tensor:
def forward(self, hidden_states: torch.Tensor, *, attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs) -> torch.Tensor:
assert position_embeddings is not None
residual = hidden_states