feat: set rlimit in main

AshishKumar4 · Aug 3, 2024 · 0b90a38 · 0b90a38
1 parent 4555209
commit 0b90a38
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/training.py b/training.py
@@ -38,6 +38,8 @@
 import cv2
 import argparse
 
+import resource
+
 #####################################################################################################################
 ################################################# Initialization ####################################################
 #####################################################################################################################
@@ -881,6 +883,14 @@ def boolean_string(s):
                     default=False, help='Run the dataset iterator for 3000 steps for testintg/benchmarking')
 
 def main(args):
+    resource.setrlimit(
+        resource.RLIMIT_CORE,
+        (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
+
+    resource.setrlimit(
+        resource.RLIMIT_OFILE,
+        (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
+
     if args.load_from_checkpoint:
         assert args.experiment_name is not None and args.experiment_name != "", "Experiment name needs to be given if load_from_checkpoint is True"