Add ids for better display.

NVIDIA · Sep 18, 2024 · 6fdcc6b · 6fdcc6b
1 parent 5ca5e11
commit 6fdcc6b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/python/test_transformer_engine.py b/tests/python/test_transformer_engine.py
@@ -22,7 +22,11 @@ class ComputeType(Enum):
 
 
 @pytest.mark.mpi
-@pytest.mark.parametrize("compute_type", [ComputeType.FORWARD, ComputeType.BACKWARD])
+@pytest.mark.parametrize(
+    "compute_type",
+    [ComputeType.FORWARD, ComputeType.BACKWARD],
+    ids=["forward", "backward"],
+)
 def test_transformer_layer(multidevice_test, benchmark, compute_type):
     # Hyperparameters for GPT-3
     hidden_size = 12288