Make lazy tensor ptr class customizable (#76476)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/76476

Test Plan: Imported from OSS

Reviewed By: Krovatkin, bdhirsh

Differential Revision: D35980433

Pulled By: wconstab

fbshipit-source-id: 1d4d00a494bf8aea86278b007f7f353cd7a822f8
(cherry picked from commit a78655bef23b5fa8487ced13443ca0bfdec65e5c)
This commit is contained in:
Will Constable
2022-04-27 20:16:46 -07:00
committed by PyTorch MergeBot
parent 4cae57080a
commit d0cb31d5bc
2 changed files with 6 additions and 3 deletions

View File

@ -279,6 +279,7 @@ def run_gen_lazy_tensor(
create_aten_from_ltc_tensor: str = "torch::lazy::CreateAtenFromLtcTensor",
tuple_aten_from_ltc_tensors: str = "torch::lazy::TupleAtenFromLtcTensors",
lazy_value_class: str = "torch::lazy::Value",
lazy_tensor_ptr: str = "LazyTensorPtr",
) -> None:
template_dir = os.path.join(aten_path, "templates")
@ -482,6 +483,7 @@ def run_gen_lazy_tensor(
create_aten_from_ltc_tensor,
tuple_aten_from_ltc_tensors,
lazy_value_class,
lazy_tensor_ptr,
),
grouped_native_functions,
codegenInplaceVariant=True,