From 0a9b7a7b215271f1aca1277687bc256e989a9b35 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Mon, 23 Dec 2024 14:31:55 -0800
Subject: [PATCH] Add machine-readable list of PyTorch GPU tests.

Used internally for benchmarking.

PiperOrigin-RevId: 709153019
---
 test/kubernetes/benchmarks/pytorch.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/test/kubernetes/benchmarks/pytorch.go b/test/kubernetes/benchmarks/pytorch.go
index 4d382fcacb..e956c13eb6 100644
--- a/test/kubernetes/benchmarks/pytorch.go
+++ b/test/kubernetes/benchmarks/pytorch.go
@@ -144,6 +144,15 @@ var (
 			test:   eval,
 		},
 	}
+
+	// AllTests is a map of test names to the tests.
+	AllTests = map[string][]pytorchTest{
+		"FastNLPBert":       FastNLPBert,
+		"BigBird":           BigBird,
+		"SpeechTransformer": SpeechTransformer,
+		"LearningToPaint":   LearningToPaint,
+		"MobileNetV2":       MobileNetV2,
+	}
 )
 
 // Name returns the name of the test with the argument parameters included. It is formatted so