Enable memory metrics in tests that need it (huggingface#11859)

NielsRogge · May 25, 2021 · 6da129c · 6da129c
1 parent db0b247
commit 6da129c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
@@ -1102,7 +1102,7 @@ def check_mem_metrics(self, trainer, check_func):
     def test_mem_metrics(self):
 
         # with mem metrics enabled
-        trainer = get_regression_trainer()
+        trainer = get_regression_trainer(skip_memory_metrics=False)
         self.check_mem_metrics(trainer, self.assertIn)
 
         # with mem metrics disabled
@@ -1123,7 +1123,7 @@ def test_fp16_full_eval(self):
         b = torch.ones(1000, bs) - 0.001
 
         # 1. with mem metrics enabled
-        trainer = get_regression_trainer(a=a, b=b, eval_len=16)
+        trainer = get_regression_trainer(a=a, b=b, eval_len=16, skip_memory_metrics=False)
         metrics = trainer.evaluate()
         del trainer
         gc.collect()
@@ -1144,7 +1144,7 @@ def test_fp16_full_eval(self):
         self.assertLess(fp32_eval, 5_000)
 
         # 2. with mem metrics disabled
-        trainer = get_regression_trainer(a=a, b=b, eval_len=16, fp16_full_eval=True)
+        trainer = get_regression_trainer(a=a, b=b, eval_len=16, fp16_full_eval=True, skip_memory_metrics=False)
         metrics = trainer.evaluate()
         fp16_init = metrics["init_mem_gpu_alloc_delta"]
         fp16_eval = metrics["eval_mem_gpu_alloc_delta"]