Hi,
thank you very much. This helps a lot.
This is the code of my training loop:
def train(train_loader, model, optimizer, criterion, n_epochs, batch_size):
    cost = []
    metric_hist = []
    for epoch in range(n_epochs):
        total_loss = total_metrics = 0
        for i, (x, y) in enumerate(train_loader):
            # reshape input to 4D   =>  [batch_size, N_channels=1, Height, Width]
            x = np.reshape(x, [batch_size, 1, x.shape[1], x.shape[2]])
            # reshape target to 2D  =>  [batch_size, N_channels=1]
            y = np.reshape(y, [batch_size, 1])

            # convert to float and load to GPU
            x = x.float().to(device)
            y = y.float().to(device)

            model.zero_grad()

            # forward
            y_pred = model(x)

            # calc loss & rmse
            loss = criterion(y_pred, y)                                     # LOSS = Mean Square Error
            metrics = rmse(y_pred, y)                                    # METRICS = Root Mean Square Error

            # backward
            optimizer.zero_grad()
            loss.backward(retain_graph=True)
            optimizer.step()

            # append results
            total_loss += loss.item()
            total_metrics += metrics.item()

            # print status every 100th step
            if (i+1) % 100 == 0:
                print(i+1, " out of ", len(train_loader), "iterations: ",
                      round(((i+1) / len(train_loader)) * 100, 2), "% Done -- LOSS: ", round(loss.item(), 3),
                      "-- RMSE: ", round(metrics.item(), 3))

        # calc average loss of epoch and append to cost
        avg_metrics = total_metrics / len(train_loader)
        avg_loss = total_loss / len(train_loader)
        cost.append(avg_loss)
        metric_hist.append(avg_metrics)

        print(epoch+1, " out of ", n_epochs, "epochs: ", round(((epoch+1) / n_epochs) * 100, 3)
              , "% Done --", "LOSS: ", round(avg_loss, 2), "-- RMSE: ", round(avg_metrics, 3), '\n')

    return cost, metric_hist