limit epochs

Change-Id: Ibfabea0833dfaf3e9f3922438c0d5dca97463592
GoogleCloudPlatform · Feb 15, 2024 · 5eff685 · 5eff685
1 parent dbb523f
commit 5eff685
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/dags/legacy_test/tests/pytorch/nightly/resnet50-mp.libsonnet b/dags/legacy_test/tests/pytorch/nightly/resnet50-mp.libsonnet
@@ -149,6 +149,7 @@ local tpus = import 'templates/tpus.libsonnet';
     command+: [
       '--pjrt_distributed',
       '--nometrics_debug',
+      '--num_epochs=2',
     ],
     flags+: {
       modelDir: null,