Hello! I want to get the validation loss after each epoch while training. I have this piece of code:
def train(model, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        scheduler.step()            
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)

        loss = torch.sqrt(F.mse_loss(output, target))
        loss.backward()
        optimizer.step()

        loss_validation = torch.sqrt(F.mse_loss(model(factors_val), product_val))

        if batch_idx%100 == 0:
            print(loss)
    return(loss, loss_validation)

I have a batch size of 512 and 300 epoch. The test size has 250000 inputs and the validation set has 20000. The NN is a simple feed forward fully connected with 8 hidden layers. If I don’t use loss_validation = torch.sqrt(F.mse_loss(model(factors_val), product_val)) the code works fine. However, if I use that line, I am getting a CUDA out of memory message after epoch 44. The way I go through the epochs is this:
for epoch in range(1, epochs ):
    losses = losses + [train(model, train_loader, optimizer)]

All I want to do is to save the training and validation set after each epoch so I can plot it. Can someone tell me what am I doing wrong and how can I fix it? I am really new to pytorch so it is probably something stupid and obvious, but I really need some help. Thank you!