@ptrblck
def eval_test(model, test_dataloader, criterion, device):
    """Evaluation model on test set"""
    loss = []
    model.eval()
    with torch.no_grad():
        for _dict in test_dataloader:
            _dict = {key:var.to(device) for key, var in _dict.items()}
            preds = model(_dict['cont'], _dict['cat'])
            loss.append(criterion(preds, _dict['target'], yr_weights=_dict['wts']).item())
    return np.mean(loss)


I believe .item() is detaching it. Also I am using aws cloud instance, could this happen due to any version compatibility? I am so puzzled by this bug.