Custom loss function using argsort

Try printing requires_grad for each tensor in your loss. For instance, if data_nd.requires_grad if False, then it won’t work.