Thanks for the tip about torch.sparse: from the docs it seems to use the COO format which should also work well 
And thanks for clarifying the reason for encoding the CSR format by hand - when I find a solution to the torch > 1.5 issue, I’ll expand the text accordingly!