From a21458f3c060e9f060981dda5317a480c1f6be24 Mon Sep 17 00:00:00 2001
From: Aditya Venkataraman <aditya.venky@gmail.com>
Date: Mon, 24 Nov 2025 18:02:44 -0800
Subject: [PATCH] Fix import path for register_op_strategy

Changed in PyTorch as part of https://github.com/pytorch/pytorch/pull/168221
---
 autoparallel/dtensor_util/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/autoparallel/dtensor_util/utils.py b/autoparallel/dtensor_util/utils.py
index 3341e2e9..09f3c11c 100644
--- a/autoparallel/dtensor_util/utils.py
+++ b/autoparallel/dtensor_util/utils.py
@@ -17,10 +17,10 @@
     OpStrategy,
     StrategyType,
 )
+from torch.distributed.tensor._ops.registration import register_op_strategy
 from torch.distributed.tensor._ops.utils import (
     generate_redistribute_costs,
     is_tensor_shardable,
-    register_op_strategy,
 )
 from torch.distributed.tensor.placement_types import Placement, Replicate, Shard