From 87dd0929264068ce2890edbe37ead9eebe70be2e Mon Sep 17 00:00:00 2001 From: Shabab Ayub Date: Wed, 23 Oct 2024 09:22:24 -0700 Subject: [PATCH] another torch deploy protection :( Summary: this import doesnt work in torch dceploy for whatever reason (we're getting rid of it all in 2 months so just bare with this for now) Differential Revision: D64830919 --- torchrec/distributed/train_pipeline/train_pipelines.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchrec/distributed/train_pipeline/train_pipelines.py b/torchrec/distributed/train_pipeline/train_pipelines.py index 85de6e3b8..27092bf11 100644 --- a/torchrec/distributed/train_pipeline/train_pipelines.py +++ b/torchrec/distributed/train_pipeline/train_pipelines.py @@ -74,7 +74,8 @@ except ImportError: logger.warning("torchrec_use_sync_collectives is not available") -torch.ops.import_module("fbgemm_gpu.sparse_ops") +if not torch._running_with_deploy(): + torch.ops.import_module("fbgemm_gpu.sparse_ops") class ModelDetachedException(Exception):