From 31793b28a75dd72157366671810b10e8fb2f7419 Mon Sep 17 00:00:00 2001 From: Jon Bolin Date: Thu, 5 Oct 2023 21:55:27 +0000 Subject: [PATCH] Don't WrapXlaData --- torch_xla/csrc/init_python_bindings.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp index fe18f9508b6..421066ba72c 100644 --- 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));