Merge pull request #31 from 36000/buffer_fixes

avoid buffer size overflow

Author: 36000

cuslines/cuda_python/cu_propagate_seeds.py

@@ -1,4 +1,5 @@
 import numpy as np
+import math
 import gc
 from cuda.bindings import runtime
 from cuda.bindings.runtime import cudaMemcpyKind
@@ -240,7 +241,7 @@ def get_buffer_size(self):
             lens = self.sline_lens[ii]
             for jj in range(self.nSlines[ii]):
                 buffer_size += lens[jj] * 3 * REAL_SIZE
-        return buffer_size
+        return math.ceil(buffer_size / MEGABYTE)
 
     def as_generator(self):
         def _yield_slines():
@@ -256,4 +257,4 @@ def _yield_slines():
         return _yield_slines()
 
     def as_array_sequence(self):
-        return ArraySequence(self.as_generator(), self.get_buffer_size() // MEGABYTE)
+        return ArraySequence(self.as_generator(), self.get_buffer_size())

cuslines/cuda_python/cu_tractography.py

@@ -248,7 +248,7 @@ def generate_sft(self, seeds, ref_img):
                     seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz].shape[0]
                 )
         array_sequence = ArraySequence(
-            (item for gen in generators for item in gen), buffer_size // MEGABYTE
+            (item for gen in generators for item in gen), buffer_size
         )
         return StatefulTractogram(array_sequence, ref_img, Space.VOX)