torchrec: '_OpNamespace' object has no attribute 'jagged_2d_to_dense'

[BorisPolonsky]

Upon setting up customized torchrec as dependency of hpcaitech/CachedEmbedding following the official README.md, torchrec does not appear to be properly installed.

Steps to replicate

1. Install PyTorch 1.12.1 build with CUDA11.6 pip3 install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116
2. Install customized torchrec git clone git clone https://github.com/hpcaitech/torchrec&& cd torchrec && pip install .

And run in python

import torchrec
libnvrtc.so.11.2: cannot open shared object file: No such file or directory
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/torch/_ops.py", line 198, in __getattr__
    op, overload_names = torch._C._jit_get_operation(qualified_op_name)
RuntimeError: No such operator fbgemm::jagged_2d_to_dense

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/polonsky/Documents/cached-embd-build/hpcaitech-torchrec/torchrec/torchrec/__init__.py", line 8, in <module>
    import torchrec.distributed  # noqa
  File "/home/polonsky/Documents/cached-embd-build/hpcaitech-torchrec/torchrec/torchrec/distributed/__init__.py", line 36, in <module>
    from torchrec.distributed.model_parallel import DistributedModelParallel  # noqa
  File "/home/polonsky/Documents/cached-embd-build/hpcaitech-torchrec/torchrec/torchrec/distributed/model_parallel.py", line 20, in <module>
    from torchrec.distributed.embedding import EmbeddingCollectionSharder
  File "/home/polonsky/Documents/cached-embd-build/hpcaitech-torchrec/torchrec/torchrec/distributed/embedding.py", line 30, in <module>
    from torchrec.distributed.embedding_sharding import (
  File "/home/polonsky/Documents/cached-embd-build/hpcaitech-torchrec/torchrec/torchrec/distributed/embedding_sharding.py", line 15, in <module>
    from torchrec.distributed.dist_data import (
  File "/home/polonsky/Documents/cached-embd-build/hpcaitech-torchrec/torchrec/torchrec/distributed/dist_data.py", line 16, in <module>
    from torchrec.distributed.comm_ops import (
  File "/home/polonsky/Documents/cached-embd-build/hpcaitech-torchrec/torchrec/torchrec/distributed/comm_ops.py", line 27, in <module>
    import fbgemm_gpu  # @manual  # noqa
  File "/usr/local/lib/python3.8/site-packages/fbgemm_gpu/__init__.py", line 22, in <module>
    from . import _fbgemm_gpu_docs
  File "/usr/local/lib/python3.8/site-packages/fbgemm_gpu/_fbgemm_gpu_docs.py", line 18, in <module>
    torch.ops.fbgemm.jagged_2d_to_dense,
  File "/usr/local/lib/python3.8/site-packages/torch/_ops.py", line 202, in __getattr__
    raise AttributeError(f"'_OpNamespace' object has no attribute '{op_name}'") from e
AttributeError: '_OpNamespace' object has no attribute 'jagged_2d_to_dense'

[GuanhuaWang]

same issue happens on my side as well