Update tests/accelerators/test_dp.py

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>
Lightning-AI · Mar 4, 2021 · 6b3d47e · 6b3d47e
1 parent 56acaf1
commit 6b3d47e
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/tests/accelerators/test_dp.py b/tests/accelerators/test_dp.py
@@ -130,13 +130,7 @@ def test_dp_training_step_dict(tmpdir):
     """
     This test verify dp properly reduce dictionaries
     """
-
-    class TestModel(BoringModel):
-
-        def training_step(self, batch, batch_idx):
-            return super().training_step(batch, batch_idx)
-
-    model = TestModel()
+    model = BoringModel()
     model.training_step_end = None
     trainer = pl.Trainer(
         default_root_dir=tmpdir,