From 4d6af1d8b96c0c9e2a998906a7fd1f2e7dc8f20b Mon Sep 17 00:00:00 2001 From: xzeng Date: Thu, 16 Mar 2023 12:44:57 -0400 Subject: [PATCH] fix tensor device issue https://github.com/nv-tlabs/LION/issues/31#issue-1627736930 the original code index cpu tensor with cuda tensor (works in torch 1.10.2), may fail in other torch version? --- utils/eval_helper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/eval_helper.py b/utils/eval_helper.py index 95585fb..d817179 100644 --- a/utils/eval_helper.py +++ b/utils/eval_helper.py @@ -52,7 +52,7 @@ def compute_NLL_metric(gen_pcs, ref_pcs, device, writer=None, output_name='', ba for k in metrics.keys(): sorted, indices = torch.sort(metrics[k]) - worse_ten, worse_score = indices[-10:], sorted[-10:] + worse_ten, worse_score = indices[-10:].cpu(), sorted[-10:].cpu() titles = 'nll/worst-%s-%s' % (k, tag) subtitles = [['ori', 'gen-%s=%.2fx1e-2' % (k, worse_score[j]*1e2)] for j in range(len(worse_score))]