Limit test_pytorch_lightning to 1 device

http://b/238238348

Author: djherbis

tests/test_pytorch_lightning.py

@@ -62,7 +62,7 @@ def test_version(self):
     def test_mnist(self):
         dm = LitDataModule()
         model = LitClassifier()
-        trainer = pl.Trainer(max_epochs=1)
+        trainer = pl.Trainer(max_epochs=1, log_every_n_steps=1, devices=1)
         trainer.fit(model, datamodule=dm)
 
         self.assertIn("train_loss", trainer.logged_metrics)