Loss não diminui

Olá, estou com o mesmo problema do post anterior. meu modelo não converge de maneira alguma ficando sempre com a mesma Loss. Não sei se está relacionado ao fato de eu não estar usando GPU agora (o google colab não deixa mais eu usar por um tempo) Eu segui os mesmos parametros que foram sugeridos acima e não funcionou. Alguma outra ideia?

args = {
    'batch_size': 100,
    'num_workers': 4,
    'num_classes': 10,
    'epochs': 10,
    'weight_decay': 0,
    'lr': 0.1
}

if torch.cuda.is_available():
  args['device'] = torch.device('cuda')
else:
  args['device'] = torch.device('cpu')

print(args)

Loader:

train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_set, batch_size=args['batch_size'], shuffle=True, num_workers=args['num_workers'])
test_loader = DataLoader(test_set, batch_size=args['batch_size'], shuffle=False, num_workers=args['num_workers'])

Modelo:

class MNISTClassifier(nn.Module):

  def __init__(self,input_size, hidden_size, output_size) -> None:
     super().__init__()
     self.features = nn.Sequential(
         nn.Linear(input_size, hidden_size),
         nn.ReLU(),
         nn.Linear(hidden_size, hidden_size),
         nn.ReLU()
     )
     self.out = nn.Linear(hidden_size, output_size)
     self.softmax = nn.Softmax()

  def forward(self, x):
    x = x.view(x.size(0), -1)
    x = self.features(x)
    x = self.softmax(self.out(x))
    return x


input_size = 28 * 28
hidden_size = 128
output_size = args['num_classes']
MNIST_model = MNISTClassifier(input_size, hidden_size, output_size).to(args['device'])

Otimizadores e perda:

#loss
criterion = nn.CrossEntropyLoss().to(args['device'])
#optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])
#

Execução

for epoch in range(args['epochs']):
  epoch_loss = []
  for batch in train_loader:
    data, target = batch
    #cast to device
    data = data.to(args['device'])
    target = target.to(args['device'])

    #forward
    pred = MNIST_model(data)
    loss = criterion(pred, target)
    epoch_loss.append(loss.cpu().data)

    #backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

  epoch_loss = np.asarray(epoch_loss)
  print("Epoca %d, Loss %.4f +/- %.4f" % (epoch, epoch_loss.mean(), epoch_loss.std()) )