Hi Thomas,
I see, thank you!
Working with preallocated tensor wouldn’t be feasible in my case as I would like the gradient to be calculated only for certain rows of that large tensor, which is also not possible in Pytorch if I understand well? It is possible that the gradient is set to 0 for the rows not used, but I need them to not be calculated - so that backward time doesn’t depend on the total number of rows in that large matrix, but only on the number of rows used in that operation (and their size ofc). Do you have an idea on how to achieve this?
I found this to be similar to torch.nn.functional.embedding, but there backward time (and also optimizer step time) depend on the total size of the values matrix.
Best regards,
Bozidar