No, this will not work. As I said, if you want your computation to be compatible with autograd, it needs to be executed on Variables from the start to the very end. You can’t unpack the tensors and repack them in the middle, because they won’t be connected to the initial graph, and will not forward the gradient. In your snippet no grad will be sent to X and y, because the backward will end on X_silimarity and association (they are graph leaves - their .creator is None).