That’s what’s puzzling me as well. The only saving of loss that I do is in storing the loss for tensorboard. Here’s my training loop
 for epoch in tqdm(range(START_EPOCH, NUM_EPOCHS)):
        if ENABLE_LOGS:
            torch.save(
                get_model_dict(model, optim, sched, epoch, cnt, cnt_val),
                'NeuralBP_featsize_{}_batchsize_{}_epoches_{}.pt'.format(NAME, OUT_FEAT_SIZE, BSIZE, epoch))
        # Training Loop
        for i, data in tqdm(enumerate(train_loader)):
            optim.zero_grad()

            # Feed-forward
            data = data.to(device)
            out = model(data)
            out = torch.softmax(out, dim=1)
            out = log_sinkhorn(out, num_it=10)

            loss = loss_fn(out.reshape(-1), data.lbl_mat.reshape(-1))
            loss.backward()
            torch.nn.utils.clip_grad_norm(parameters, 1.0)
            optim.step()
            loss_val = float(loss)
            if ENABLE_LOGS:
                writer.add_scalar('loss', loss_val, cnt)
            cnt += 1
        sched.step()
        # Validation loop
        if epoch % VAL_EPOCH == 0:
            with torch.no_grad():
                for j, data in tqdm(enumerate(val_loader)):
                    data = data.to(device)
                    out = model.forward(data)

                    res = torch.softmax(out, dim=1)
                    res = log_sinkhorn(res, num_it=10)
                    tmax = torch.max(res, dim=1, keepdim=True)[0]
                    mask = res.ge(tmax)
                    res[mask] = 1.0
                    res[~mask] = 0.0

                    acc = accuracy_score(res.cpu().numpy().reshape(-1), data.lbl_mat.cpu().numpy().reshape(-1))
                    prec = precision_score(res.cpu().numpy().reshape(-1), data.lbl_mat.cpu().numpy().reshape(-1))
                    recall = recall_score(res.cpu().numpy().reshape(-1), data.lbl_mat.cpu().numpy().reshape(-1))
                    f1 = f1_score(res.cpu().numpy().reshape(-1), data.lbl_mat.cpu().numpy().reshape(-1))
                    if ENABLE_LOGS:
                        writer.add_scalar('acc', acc, cnt_train)
                        writer.add_scalar('prec', prec, cnt_train)
                        writer.add_scalar('recall', recall, cnt_train)
                        writer.add_scalar('f1', f1, cnt_train)

                    cnt_train += 1