Just a bit of cleanup/flexibility/catching another warning (that it raised when I used it with dask).