미니배치를 이용하여 강의 내용과 같이 만들어서 cost plot을 그려보았지만 수렴하지 않는 것 같습니다.

혹시 제 코드에 이상한 점이 있거나, 이를 해결할 수 있는 방법을 알려주실 수 있을까요?

감사합니다.

import torch

class MultivariateLinearRegressionModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 1)

    def forward(self, x):
        return self.linear(x)

class CustomDataset(torch.utils.data.Dataset):
  def __init__(self):
    self.x_data = [[73, 80, 75],
                   [93, 88, 93],
                   [89, 91, 80],
                   [96, 98, 100],
                   [73, 66, 70]]
    self.y_data = [[152], [185], [180], [196], [142]]


  def __len__(self):
    return len(self.x_data)

  def __getitem__(self, idx):
    x = torch.FloatTensor(self.x_data[idx])
    y = torch.FloatTensor(self.y_data[idx])

    return x, y
    
model = MultivariateLinearRegressionModel()
dataset = CustomDataset()
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True)
optimizer = torch.optim.SGD(model.parameters(), lr = 1e-5)
epochs = 20
cost_list = []

for epoch in range(0, epochs):
  for batch_idx, samples in enumerate(dataloader):
    x_train, y_train = samples
    hypothesis = model(x_train)
    cost = torch.nn.functional.mse_loss(hypothesis, y_train)
    optimizer.zero_grad()
    cost.backward()
    optimizer.step()
    cost_list.append(cost.item())

    print('Epoch {:4d}/{} Batch:{}/{} Cost: {:.6f}'.format(
        epoch, epochs, batch_idx+1, len(dataloader), cost.item()))

import matplotlib.pyplot as plt
print(cost_list)
print(len(cost_list))
plt.plot(cost_list)
plt.ylim([0, 100])
plt.show()