I assume you’ve transformed your parameters to the mkl layout before storing the state_dict?
If so, could you try to load the state_dict and transform the parameters back to a dense representation via:
state_dict = torch.load(...)
for key in state_dict:
    state_dict[key] = state_dict[key].to_dense()
model = MyModel(...)
model.load_state_dict(state_dict)

Could you also please create an issue?