Solved by InnovArul in post #2 
              
              
                Can you try enclosing the above code under with torch.no_grad():? 
model.eval()
with torch.no_grad():
    val_loss_epoch = valid(vdl, model, criterion, writer)