Could Torchrec support PyTorch's PrivateUse1 Dispatch Key?

[kwgqjj]

Hello,

I've noticed that there are many conditional checks like if device.type == "cuda" in our TorchRec codebase. Without modifying TorchRec's source code, such fixed conditional logic might not be flexible enough to conveniently support third-party devices. From what I understand, PyTorch has introduced the PrivateUse1 DispatchKey to address third-party device extension issues. I'd like to ask if our TorchRec repository could add support for PyTorch's PrivateUse1 DispatchKey? This would enable third-party devices to seamlessly adapt TorchRec's functionality through PrivateUse1 without requiring code modifications.