Sparse CSR tensors crash `state_dict()` (1.11 nightly)

[Linux-cpp-lisp]

🐛 Describe the bug

This code works:

import torch

class SparseTensorModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("a", torch.eye(3).to_sparse())

    def forward(self):
        pass

s = SparseTensorModule()
s.a
print("hi")
print(s.state_dict())

producing:

hi
OrderedDict([('a', tensor(indices=tensor([[0, 1, 2],
                       [0, 1, 2]]),
       values=tensor([1., 1., 1.]),
       size=(3, 3), nnz=3, layout=torch.sparse_coo))])

but modifying

self.register_buffer("a", torch.eye(3).to_sparse())

to

self.register_buffer("a", torch.eye(3).to_sparse_csr())

gives

hi
Segmentation fault (core dumped)

Versions

PyTorch version 1.11.0.dev20220119
CUDA available
Ubuntu Linux 21.10

cc @nikitaved @pearu @cpuhrsch

[cpuhrsch]

I'd just like to add that I can reproduce this error on my end and am looking into it.

[cpuhrsch]

This should be resolved now that #71963 landed.