How to put tensor on a custom Function to cuda device?

[zetyquickly (Emil Bogomolov)]

Here is an answer, that might work for you.
You should create Tensor from NumPy array and then transfer it to device

Converting numpy array to tensor on GPU

You should transform numpy arrays to PyTorch tensors with torch.from_numpy. Otherwise some weird issues might occur. img = torch.from_numpy(img).float().to(device)