diff --git a/torchrec/distributed/tests/test_infer_hetero_shardings.py b/torchrec/distributed/tests/test_infer_hetero_shardings.py
index d101121a3..f7db81c1f 100755
--- a/torchrec/distributed/tests/test_infer_hetero_shardings.py
+++ b/torchrec/distributed/tests/test_infer_hetero_shardings.py
@@ -15,7 +15,6 @@
 import hypothesis.strategies as st
 
 import torch
-from caffe2.torch.fb.model_transform.splitting.split_dispatcher import SplitDispatchMode
 from hypothesis import given, settings
 from torchrec import EmbeddingBagConfig, EmbeddingCollection, EmbeddingConfig
 from torchrec.distributed.embedding_types import ShardingType