uncomment memory checking

eth-easl · Jan 12, 2024 · 490bc83 · 490bc83
1 parent fac49b7
commit 490bc83
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/nanotron/trainer.py b/src/nanotron/trainer.py
@@ -170,8 +170,8 @@ def __init__(self, config_or_config_file: Union[Config, str]):
             group=self.parallel_context.world_pg,
             rank=None,
         )
-        # if free_mem < MIN_GPU_MEM_THRESHOLD:
-        #     raise RuntimeError(f"Not enough memory to train the model on node {os.environ.get('SLURMD_NODENAME')}")
+        if free_mem < MIN_GPU_MEM_THRESHOLD:
+            raise RuntimeError(f"Not enough memory to train the model on node {os.environ.get('SLURMD_NODENAME')}")
         # Try to allocate all the memory
         test_tensor_size = int(free_mem * 0.9)
         test_tensor = torch.zeros((test_tensor_size,), dtype=torch.uint8, device=torch.device("cuda"))