Move early_stop_checker GPU test to dedicated file (#763)

Summary: Pull Request resolved: #763

Reviewed By: anshulverma

Differential Revision: D55388225

fbshipit-source-id: 69a37acdb2098ff1b3922e1b98792df8e4ec82a3

Author: diego-urgell

tests/utils/test_early_stop_checker.py

@@ -11,7 +11,6 @@
 
 import torch
 from torchtnt.utils.early_stop_checker import EarlyStopChecker
-from torchtnt.utils.test_utils import skip_if_not_gpu
 
 
 class EarlyStopCheckerTest(unittest.TestCase):
@@ -86,38 +85,6 @@ def test_early_stop_min_delta(self) -> None:
         should_stop = es2.check(torch.tensor(0.26))
         self.assertTrue(should_stop)
 
-    @skip_if_not_gpu
-    def test_early_stop_min_delta_on_gpu(self) -> None:
-        device = torch.device("cuda:0")
-
-        # Loss decreases beyond 0.25 but not more than min_delta
-        losses = [
-            torch.tensor([0.4], device=device),
-            torch.tensor([0.38], device=device),
-            torch.tensor([0.31], device=device),
-            torch.tensor([0.25], device=device),
-            torch.tensor([0.27], device=device),
-            torch.tensor([0.24], device=device),
-        ]
-        es1 = EarlyStopChecker("min", 3, min_delta=0.05)
-        es2 = EarlyStopChecker("min", 4, min_delta=0.05)
-
-        for loss in losses:
-            should_stop = es1.check(torch.tensor(loss))
-            self.assertFalse(should_stop)
-            should_stop = es2.check(torch.tensor(loss))
-            self.assertFalse(should_stop)
-
-        # Patience should run out
-        should_stop = es1.check(torch.tensor(0.25))
-        self.assertTrue(should_stop)
-
-        # es2 has more patience than es1
-        should_stop = es2.check(torch.tensor(0.25))
-        self.assertFalse(should_stop)
-        should_stop = es2.check(torch.tensor(0.26))
-        self.assertTrue(should_stop)
-
     def test_early_stop_max_mode(self) -> None:
 
         # Loss increases beyond 0.38 but not more than min_delta

tests/utils/test_early_stop_checker_gpu.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import unittest
+
+import torch
+from torchtnt.utils.early_stop_checker import EarlyStopChecker
+from torchtnt.utils.test_utils import skip_if_not_gpu
+
+
+class EarlyStopCheckerGPUTest(unittest.TestCase):
+    @skip_if_not_gpu
+    def test_early_stop_min_delta_on_gpu(self) -> None:
+        device = torch.device("cuda:0")
+
+        # Loss decreases beyond 0.25 but not more than min_delta
+        losses = [
+            torch.tensor([0.4], device=device),
+            torch.tensor([0.38], device=device),
+            torch.tensor([0.31], device=device),
+            torch.tensor([0.25], device=device),
+            torch.tensor([0.27], device=device),
+            torch.tensor([0.24], device=device),
+        ]
+        es1 = EarlyStopChecker("min", 3, min_delta=0.05)
+        es2 = EarlyStopChecker("min", 4, min_delta=0.05)
+
+        for loss in losses:
+            should_stop = es1.check(torch.tensor(loss))
+            self.assertFalse(should_stop)
+            should_stop = es2.check(torch.tensor(loss))
+            self.assertFalse(should_stop)
+
+        # Patience should run out
+        should_stop = es1.check(torch.tensor(0.25))
+        self.assertTrue(should_stop)
+
+        # es2 has more patience than es1
+        should_stop = es2.check(torch.tensor(0.25))
+        self.assertFalse(should_stop)
+        should_stop = es2.check(torch.tensor(0.26))
+        self.assertTrue(should_stop)