Update optical_flow.py

[bmad4ever]

Dear Maintainer,

This minor change should enable the use of single-channel images in remap without breaking or changing the implemented functionalities.

Prior to this request, I tried using expand_dims in the input image passed to image_transform_optical_flow, but the image sent to center_crop_image would not keep the 3rd shape value (I did not investigate why since this alternative solution was trivial to implement).

Best Regards