From fde9f6450205b8b26cce2aa5e9c581b66257e9bb Mon Sep 17 00:00:00 2001
From: Shuai Yang <shuaiyang@meta.com>
Date: Tue, 21 May 2024 16:21:40 -0700
Subject: [PATCH] Import sparse_ops

Differential Revision: D57648027
---
 torchrec/sparse/jagged_tensor.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/torchrec/sparse/jagged_tensor.py b/torchrec/sparse/jagged_tensor.py
index 440a1753b..5c0f22e47 100644
--- a/torchrec/sparse/jagged_tensor.py
+++ b/torchrec/sparse/jagged_tensor.py
@@ -13,6 +13,8 @@
 
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
+import fbgemm_gpu.sparse_ops  # noqa: F401, E402
+
 import torch
 from torch.autograd.profiler import record_function
 from torch.fx._pytree import register_pytree_flatten_spec, TreeSpec