/black

/black

BTW, as the numpy version unravel_index is slightly faster than PyTorch version, if running the fast training tutorial with numpy data, the total time can be 1s faster.

Thanks.

/black

Now the training speed is same as the numpy version benchmark of 0.7 release

So with the data on the GPU, we're only as fast as the numpy implementation with all on the CPU?

BTW, as the numpy version unravel_index is slightly faster than PyTorch version

We could change the logic to only use torch if the data is already on the GPU. If on the CPU, use numpy regardless of whether input was torch or numpy:

if isinstance(x, torch.Tensor) and x.device is not torch.device("cpu"):
    torch.unravel
else:
    np.unravel

Hi @rijobro ,

Thanks for your review.
There are 2 data in this case: the image data and the bg / fg indices data.
If we move the image to GPU, it can help accelerate the training.
But if we put the bg / fg indices on GPU, actually it doesn't help, because when we use the indices to randomly crop image, we need to put the indices value back to CPU by torch.item(). That's why here I think we should always save the indices on CPU.
And about the unravel_index question, I would suggest to keep the logic simple and straightforward before we totally completed backends for all the transforms, just use torch APIs for tensor and numpy APIs for numpy array. Now it's just a slight difference, maybe next PyTorch version will improve it.

What do you think?

Thanks.

@Nic-Ma sounds good, thanks for the explanations!