IO consumption in flow warp

[NK-CS-ZZL]

In the lines 32-33 of 'mmedit/models/common/flow_warp.py'
, the function "flow_warp" tries to allocate a piece of memory for a 2d grid tensor.
However, this tensor is built in common memory and immediately transferred to GPU. This operation is time-consuming and I suggest directly declaring the tensor on GPU as follows:
grid_y, grid_x = torch.meshgrid(torch.arange(0, h, device=x.device, dtype=x.dtype), torch.arange(0, w, device=x.device, dtype=x.dtype)) .
In this way, the running time is reduced from 145.55ms to 0.88ms (the input feature size is 1x13x128x108x60)

And here is an example of warping.

[wangruohui]

This is actually a better implementation. You can make a PR for this improment. Thank you in advance!

One more important thing: can you make sure doing so results in a bit-wise identical result? GPU / CPU implementation usually differes a little. This function is iteratively called in the model. Some very tiny differences may affect final accuracy as they will accumulate millions of times.