AxisError with numpy 1.18

[keceli]

I got the following error, when I tried to run FFN training with numpy 1.18:
numpy.AxisError: axis 4 is out of bounds for array of dimension 4

Looks like it is due to a recent change in numpy.expand_dims. See https://numpy.org/doc/stable/reference/generated/numpy.expand_dims.html.

Here is a patch to fix the error:

diff --git a/ffn/training/inputs.py b/ffn/training/inputs.py
index d1b5c31..b9a9c7e 100644
--- a/ffn/training/inputs.py
+++ b/ffn/training/inputs.py
@@ -152,7 +152,7 @@ def load_from_numpylike(coordinates, volume_names, shape, volume_map,
     if data.ndim == 4:
       data = np.rollaxis(data, 0, start=4)
     else:
-      data = np.expand_dims(data, 4)
+      data = np.expand_dims(data, data.ndim)

     # Add flat batch dim and return.
     data = np.expand_dims(data, 0)
--