Skip to content

Commit

Permalink
another torch deploy protection :(
Browse files Browse the repository at this point in the history
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
  • Loading branch information
s4ayub authored and facebook-github-bot committed Oct 23, 2024
1 parent 1a57ce1 commit 87dd092
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion torchrec/distributed/train_pipeline/train_pipelines.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 87dd092

Please sign in to comment.