added flags to test

Signed-off-by: mikail <mkhona@nvidia.com>

Author: mkhona-nvidia

tests/ci/L0_Tests_GPU.sh

@@ -22,5 +22,5 @@ coverage run -p --source=emerging_optimizers tests/soap_smoke_test.py
 coverage run -p --source=emerging_optimizers tests/soap_mnist_test.py
 coverage run -p --source=emerging_optimizers tests/test_scalar_optimizers.py --device=cuda
 coverage run -p --source=emerging_optimizers tests/test_spectral_clipping_utils.py
-coverage run -p --source=emerging_optimizers tests/test_normalized_optimizer.py
+coverage run -p --source=emerging_optimizers tests/test_normalized_optimizer.py --device=cuda
 coverage run -p --source=emerging_optimizers tests/normalized_optimizer_convergence_test.py --device=cuda

tests/test_normalized_optimizer.py

@@ -14,11 +14,18 @@
 # limitations under the License.
 
 import torch
+from absl import flags
 from absl.testing import absltest, parameterized
 
 from emerging_optimizers.riemannian_optimizers.normalized_optimizer import ObliqueAdam, ObliqueSGD
 
 
+# Define command line flags
+flags.DEFINE_string("device", "cpu", "Device to run tests on: 'cpu' or 'cuda'")
+
+FLAGS = flags.FLAGS
+
+
 class NormalizedOptimizerFunctionalTest(parameterized.TestCase):
     """Tests for ObliqueSGD and ObliqueAdam optimizers that preserve row/column norms."""
 
@@ -29,7 +36,7 @@ def setUp(self):
         # Set seed for CUDA if available
         if torch.cuda.is_available():
             torch.cuda.manual_seed_all(1234)
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = FLAGS.device
 
     @parameterized.parameters(
         (0),