From 5eff68573ab4ee20c6f79829924c96b762fc966c Mon Sep 17 00:00:00 2001
From: Will Cromar <wcromar@google.com>
Date: Thu, 15 Feb 2024 22:54:06 +0000
Subject: [PATCH] limit epochs

Change-Id: Ibfabea0833dfaf3e9f3922438c0d5dca97463592
---
 dags/legacy_test/tests/pytorch/nightly/resnet50-mp.libsonnet | 1 +
 1 file changed, 1 insertion(+)

diff --git a/dags/legacy_test/tests/pytorch/nightly/resnet50-mp.libsonnet b/dags/legacy_test/tests/pytorch/nightly/resnet50-mp.libsonnet
index cd496fc2..0a54ea4d 100644
--- 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,