import torch
import torch.nn as nn

# 定义Transformer模型
class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, num_layers):
        super(Transformer, self).__init__()

        # 编码器和解码器的初始化
        self.encoder = nn.TransformerEncoderLayer(input_dim, hidden_dim, num_layers)
        self.decoder = nn.TransformerDecoderLayer(output_dim, hidden_dim, num_layers)

    def forward(self, src, tgt):
        # 编码器的前向传播
        enc_output = self.encoder(src)

        # 解码器的前向传播
        dec_output = self.decoder(tgt, enc_output)

        return dec_output

# 创建Transformer模型实例
input_dim = 100
output_dim = 200
hidden_dim = 256
num_layers = 4
model = Transformer(input_dim, output_dim, hidden_dim, num_layers)

# 定义输入和目标数据
src = torch.randn(50, input_dim)
tgt = torch.randn(60, output_dim)

# 进行前向传播
output = model(src, tgt)