Thanks for the feedback @tqchen @jwfromm. I’ll move the the code to the namespace tvm.utils.data, and set batch_size and num_batches through the @property decorator.
I do agree that future support of zero copy through DLPack is interesting, so it’s worth considering using tvm.runtime.ndarrays instead of numpy arrays. One question I have about this, though, is whether we should store labels as tvm.runtime.ndarrays as well as the data.
If I provide a tvm.runtime.ndarray as input to a graph runtime module (or one of the other ways to run a relay module), is the output also a tvm.runtime.ndarray?
I want to make sure that the datatype of f(data) matches the datatype of the labels so users can directly compare them.