Testing mnist inference API