Проблема с loss

В процессе обучения не падает loss. Делаю все это впервые, как я понял, на каждой эпохе loss должен снижаться. Проблема в том, что этого не происходит. Он то снижается, то резкий скачок вверх. Пробовал менять learning_rate, batch_size, перемешивал наборы данных, результат плюс-минус такой же. Может, не достаточно количества эпох?

Training Beginning

Epoch (1/5)
----------------------------------
Test Accuracy : 83.84037929672067  Test Loss : 0.6498197642719652

Epoch (2/5)
----------------------------------
Test Accuracy : 81.58830501777953  Test Loss : 0.6336045114132832

Epoch (3/5)
----------------------------------
Test Accuracy : 78.42749901224812  Test Loss : 0.7707894688222822

Epoch (4/5)
----------------------------------
Test Accuracy : 81.03516396681154  Test Loss : 0.6860072239646762

Epoch (5/5)
----------------------------------
Test Accuracy : 77.79533781114183  Test Loss : 0.7758270732061131

Training Complete

Model.py

class DeePixBiS(nn.Module):
    def __init__(self, pretrained=True):
        super().__init__()
        dense = models.densenet161(pretrained=pretrained)
        features = list(dense.features.children())
        self.enc = nn.Sequential(*features[:8])
        self.dec = nn.Conv2d(384, 1, kernel_size=1, stride=1, padding=0)
        self.linear = nn.Linear(14 * 14, 1)

    def forward(self, x):
        enc = self.enc(x)
        dec = self.dec(enc)
        out_map = F.sigmoid(dec)
        # print(out_map.shape)
        out = self.linear(out_map.view(-1, 14 * 14))
        out = F.sigmoid(out)
        out = torch.flatten(out)
        return out_map, out