Thank you for your help 



 albanD:

If you know how to write the backward for that part though, you can create a custom Function and define the backward yourself.


Yeah I can try with that. For more reference, I am currently trying to take in some tensors and with the help of those trying to transform the image. I haven’t written backward ever and have a vague intuition this won’t work, but can I use the tensors to transform the Image–then compare the image and the label with a loss and return loss.backward() in backward ? It would still break the link between the input tensors and the Image right?