Use int64 to calculate the input tensor size when splitting acquisition functions. Otherwise, we overflow for large tensors, e.g. of shape [30000, 1000, 100].
It is tricky to add a unit test to cover the limit, as we run into memory issues. Using mocking or sparse-tensors gets a bit messy, so skipping unit testing for this simple change.
Fully backwards compatible: yes
PR checklist
[X] The quality checks are all passing
[ ] The bug case / new feature is covered by tests
[ ] Any new features are well-documented (in docstrings or notebooks)
Related issue(s)/PRs: None
Summary
Use
int64to calculate the input tensor size when splitting acquisition functions. Otherwise, we overflow for large tensors, e.g. of shape[30000, 1000, 100].It is tricky to add a unit test to cover the limit, as we run into memory issues. Using mocking or sparse-tensors gets a bit messy, so skipping unit testing for this simple change.
Fully backwards compatible: yes
PR checklist