Follow-up is also tracked here. Let me know if transforming the data to tensors would work for your use case.