Implementei uma função única forward() para centralizar os fluxos de treinamento e validação. O parâmetro mode controla o comportamento da rede: quando recebe "train", a função coloca o modelo em modo de treino com net.train(), calcula a predição, a loss, executa o backward() e atualiza os pesos com o otimizador. Quando recebe "test", a função coloca o modelo em modo de avaliação com net.eval(), calcula apenas as predições, a loss e a acurácia, sem realizar atualização dos pesos.
Também substituí as chamadas antigas train(train_loader, net, epoch) e test(test_loader, net, epoch) por chamadas à nova função forward(), passando o respectivo modo de execução. Dessa forma, o código ficou mais organizado, reduziu duplicação e deixou mais clara a diferença entre o fluxo de treino e o fluxo de validação.
Segue o link: https://github.com/Moquiuti/Rede-Neural-Deep-Learning-com-PyTorch/blob/main/Carregamento_de_Dados.ipynb
Conte com o apoio da comunidade Alura na sua jornada. Abraços e bons estudos!