MemoryDataLayer now accepts dynamic batch_size