Fix tensor printing in trainer.test() (#5138)

* Fix showing test results for tensors * Fix docs * Update pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py Co-authored-by: Nicki Skafte <skaftenicki@gmail.com> * Fix lint issues Co-authored-by: Nicki Skafte <skaftenicki@gmail.com>
Lightning-AI · Feb 4, 2021 · 6c96db0 · 6c96db0
1 parent 7d33abe
commit 6c96db0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/docs/source/starter/introduction_guide.rst b/docs/source/starter/introduction_guide.rst
@@ -768,7 +768,7 @@ Once you train your model simply call ``.test()``.
 
  --------------------------------------------------------------
  TEST RESULTS
- {'test_loss': tensor(1.1703, device='cuda:0')}
+ {'test_loss': 1.1703}
  --------------------------------------------------------------
 
 You can also run the test from a saved lightning model

diff --git a/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py b/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py
@@ -301,7 +301,8 @@ def get_evaluate_epoch_results(self):
  print('-' * 80)
  for result_idx, results in enumerate(self.eval_loop_results):
  print(f'DATALOADER:{result_idx} TEST RESULTS')
- pprint(results)
+ pprint({k: (v.item() if v.numel() == 1 else v.tolist()) if isinstance(v, torch.Tensor) else v
+ for k, v in results.items()})
  print('-' * 80)
 
  results = self.eval_loop_results