Fix several long running tests in test_auto_unit

Summary: These tests were accidentally made to run longer in #463

Differential Revision: D47656034

fbshipit-source-id: fb11998c0bb500d3e3fada71a09d69d06224963b

tests/framework/test_auto_unit.py

@@ -223,8 +223,10 @@ def forward(self, x):
         input_dim = 2
         dataset_len = 10
         batch_size = 2
+        max_epochs = 2
+
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
-        train(auto_unit, dataloader)
+        train(auto_unit, dataloader, max_epochs=max_epochs)
 
         orig_module = auto_unit.module
         swa_module = auto_unit.swa_model
@@ -634,8 +636,10 @@ def forward(self, x):
         input_dim = 2
         dataset_len = 10
         batch_size = 2
+        max_epochs = 2
+
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
-        train(auto_unit, dataloader)
+        train(auto_unit, dataloader, max_epochs=max_epochs)
 
         orig_module = auto_unit.module.module
         swa_module = auto_unit.swa_model.module.module
@@ -695,9 +699,10 @@ def forward(self, x):
         input_dim = 2
         dataset_len = 10
         batch_size = 2
+        max_epochs = 1
+
         dataloader = generate_random_dataloader(dataset_len, input_dim, batch_size)
-        state = get_dummy_train_state(dataloader)
-        train(state, auto_unit)
+        train(auto_unit, dataloader, max_epochs=max_epochs)
 
         tc = unittest.TestCase()
         tc.assertTrue(custom_noop_hook_called)