Remove call to non-existing method get_batch_axis(). Test included