Right, I see! Thanks for the clarification.
Slightly off-topic question then - inside a training loss, I need to access the values of a tensor [y_true] by indices. The other tensor [y_pred] which consists of the indices, is of type float and has float values. Since I need to compute the gradient, is there any way to access values of y_true, without rounding y_pred (would like to avoid this due to its zero gradient output almost everywhere) and then doing the type conversion of it to long? Please note that it is not possible to do any interpolation in y_true in this context. A minimal example of the said loss function is as following -
def trainloss(y_pred, y_true): 
    #y_pred is of shape [100,2], y_true is of shape [64,64]
    idx = torch.round(y_pred)  
    idx = idx.long()
    loss = y_true[idx[:,0],idx[:,1]]
    loss = torch.max(loss)
return loss