Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from .usta_multi_head_attention import UstaMultiHeadAttention | |
| from .usta_layer_norm import UstaLayerNorm | |
| from .usta_mlp import UstaMLP | |
| class UstaDecoderBlock(nn.Module): | |
| def __init__(self,embedding_dim,num_heads,context_length,device): | |
| super().__init__() | |
| self.self_attention = UstaMultiHeadAttention( | |
| embedding_dim, | |
| embedding_dim, | |
| context_length, | |
| num_heads, | |
| dropout_rate=0.5, | |
| device=device | |
| ) | |
| self.norm1 = UstaLayerNorm(embedding_dim,device=device) | |
| self.mlp = UstaMLP(embedding_dim,embedding_dim,device=device) | |
| self.norm2 = UstaLayerNorm(embedding_dim,device=device) | |
| def forward(self,x): | |
| res = self.norm1(x) | |
| x = self.self_attention(x) | |
| x = self.norm1(x) | |
| x = x + res | |
| res = self.norm2(x) | |
| x = self.mlp(x) | |
| x = self.norm2(x) | |
| x = x + res | |
| return x | |