Add comment

secondmind-labs · Nov 17, 2023 · c57b45b · c57b45b
1 parent cfbb790
commit c57b45b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/trieste/acquisition/utils.py b/trieste/acquisition/utils.py
@@ -48,6 +48,7 @@ def wrapper(x: TensorType) -> TensorType:
         if length == 0:
             return fn(x)
 
+        # Use int64 to calculate the input tensor size, otherwise we can overflow for large tensors.
         elements_per_block = tf.size(x, out_type=tf.int64) / length
         blocks_per_batch = tf.cast(tf.math.ceil(split_size / elements_per_block), tf.int32)