Fix variable batch size issue