Don't WrapXlaData

pytorch · Oct 5, 2023 · 31793b2 · 31793b2
1 parent e4e7ffd
commit 31793b2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -1720,8 +1720,8 @@ void InitXlaModuleBindings(py::module m) {
               << " vs " << expected_shard_shape;
         }
 
-        auto data_handle = WrapXlaData(ShardingUtil::CreateShardedData(
-            shards, local_devices, sharding_spec));
+        auto data_handle = ShardingUtil::CreateShardedData(
+            shards, local_devices, sharding_spec);
         XLATensorPtr xla_tensor = XLATensor::Create(std::move(data_handle));
         xla_tensor->SetShardingSpec(*sharding_spec);
         auto tensor = bridge::AtenFromXlaTensor(std::move(xla_tensor));