Adding some doc

qfactorjax/qfactor_sample_jax.py

@@ -347,6 +347,12 @@ def safe_call_jited_vmaped_state_sample_sweep(
             untrys: Array,
             training_states_kets: Array,
     ) -> tuple[Array, Array[float], Array[float], Array[int], Array[bool]]:
+        """
+        We couldn't find a way to check if we are going to allocate more than
+        the GPU memory, so we created this "safe" function that calls
+        qfactor-sample and then if OOM exception is caught it recursively
+        calls qfactor-sample with half the multistarts
+        """
 
         try:
             results = _jited_loop_vmaped_state_sample_sweep(