test

ecmwf · Aug 2, 2024 · cac927d · cac927d
1 parent 118550f
commit cac927d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/diagnostics/test_checkpoint.py b/tests/diagnostics/test_checkpoint.py
@@ -15,6 +15,8 @@
 from anemoi.training.diagnostics.callbacks import AnemoiCheckpoint
 from anemoi.training.utils.jsonify import map_config_to_primitives
 
+torch.mps.set_per_process_memory_fraction(0.0)
+
 
 class DummyModel(torch.nn.Module):
     """Dummy pytorch model for testing."""