Add machine-readable list of PyTorch GPU tests.

Used internally for benchmarking. PiperOrigin-RevId: 709153019
google · Dec 23, 2024 · 0a9b7a7 · 0a9b7a7
1 parent 1a9f8b1
commit 0a9b7a7
Showing 1 changed file with 9 additions and 0 deletions.
diff --git 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