nn.transformer 결과값이 생각한대로 안나옵니다.

import torch
import torch.nn as nn
import matplotlib.pyplot as plt

class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.embedding = nn.Embedding(10, 20)  #수정
    self.transformer = nn.Transformer(20, 5,dropout = 0) #수정
  def forward(self, src, tgt):
    tgt_mask=self.transformer.generate_square_subsequent_mask(tgt.shape[1])
    src =  self.embedding(src)
    src=src.permute(1,0,2)
    tgt = tgt.permute(1,0,2)
    out = self.transformer(src, tgt, tgt_mask=tgt_mask)
    out = out.permute(1,0, 2)
    return out

critic = nn.L1Loss()
model = Net()
optim = torch.optim.Adam(model.parameters(), 3e-4)

src = torch.randint(1, 10, (50, 10)) #수정

tgt = torch.rand((50, 10, 20))#수정
tgt_input = torch.cat((torch.zeros((50, 1, 20)), tgt[:, :-1, :]), 1)  #수정

epoch = 0

model.train()
for epoch in range(10000):  #계속 돌리기 위해 수정
  optim.zero_grad()
  out=model(src, tgt_input)
  loss=critic(out, tgt)
  
  loss.backward()
  optim.step()
  if epoch%100 == 0:
    fig, (ax0, ax1)=plt.subplots(2, 1)
    ax0.pcolor(out[0].to('cpu').detach().numpy())
    ax1.pcolor(tgt[0].to('cpu').detach().numpy())
    plt.show()

  epoch +=1

model.eval()
tgt_input = torch.zeros((50, 1, 20))
with torch.no_grad():
  for _ in range(10):
    out=model(src, tgt_input)
    tgt_input = torch.cat((tgt_input, out[:, -1:, :]), 1)
  
  fig, (ax0, ax1)=plt.subplots(2, 1)
  ax0.pcolor(out[0].to('cpu').detach().numpy())
  ax1.pcolor(tgt[0].to('cpu').detach().numpy())
  plt.show()

train시 마지막 결과 그림은 다음과 같습니다.
image

eval시 결과그림은 다음과 같습니다.
image

다음과 같이 eval시 trian시의 결과값처럼 도출이 안되는데 되게하는 방법이 없을까요?

학습과 동일한 입력를 넣으셨다면 동일한 결과를 얻으실 수 있을 것 같습니다.
model.eval()
tgt_input = torch.cat((torch.zeros((50, 1, 20)), tgt[:, :-1, :]), 1) #수정
with torch.no_grad():
out=model(src, tgt_input)
fig, (ax0, ax1)=plt.subplots(2, 1)
ax0.pcolor(out[0].to(‘cpu’).detach().numpy())
ax1.pcolor(tgt[0].to(‘cpu’).detach().numpy())
plt.show()