From 6d9c4d63667a48fabf63d5caa979548113d447a8 Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Thu, 16 May 2024 11:26:10 -0700 Subject: [PATCH] Fix TBE cache operator not found issue (#2597) Summary: - Fix TBE cache operator not found issue Differential Revision: D57447659 --- fbgemm_gpu/fbgemm_gpu/sparse_ops.py | 3 ++- .../fbgemm_gpu/split_table_batched_embeddings_ops_training.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/fbgemm_gpu/fbgemm_gpu/sparse_ops.py b/fbgemm_gpu/fbgemm_gpu/sparse_ops.py index c9b0c7f37b..c181064152 100644 --- a/fbgemm_gpu/fbgemm_gpu/sparse_ops.py +++ b/fbgemm_gpu/fbgemm_gpu/sparse_ops.py @@ -13,7 +13,7 @@ from fbgemm_gpu.split_embedding_configs import SparseType from fbgemm_gpu.split_table_batched_embeddings_ops_common import PoolingMode -from torch.fx.experimental.symbolic_shapes import guard_size_oblivious + try: # pyre-ignore @@ -49,6 +49,7 @@ import torch.utils._pytree as pytree from torch import SymInt, Tensor +from torch.fx.experimental.symbolic_shapes import guard_size_oblivious if hasattr(torch.library, "register_fake"): diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py index 8c84e119ce..1e2210b0c7 100644 --- a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py +++ b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py @@ -25,6 +25,7 @@ from torch import nn, Tensor # usort:skip import fbgemm_gpu.split_embedding_codegen_lookup_invokers as invokers +import fbgemm_gpu.tbe.cache # noqa: F401 from fbgemm_gpu.runtime_monitor import ( AsyncSeriesTimer, TBEStatsReporter,