From 50f012ef924e6a940e24010f7847daddcfed3452 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Thu, 24 Jun 2021 14:41:23 -0700 Subject: [PATCH] fix(aten::slice): Allow slicing of pytorch tensors We could in the future allow evaluator / converter overwriting to do the slicing at compile time, but for now we are just going to freeze the tensor Signed-off-by: Naren Dasan Signed-off-by: Naren Dasan --- core/conversion/converters/impl/select.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/conversion/converters/impl/select.cpp b/core/conversion/converters/impl/select.cpp index 0d65a54e86..616c9907e6 100644 --- a/core/conversion/converters/impl/select.cpp +++ b/core/conversion/converters/impl/select.cpp @@ -199,7 +199,7 @@ auto select_registrations TRTORCH_UNUSED = .pattern( {"aten::slice.Tensor(Tensor(a) self, int dim=0, int start=0, int end=9223372036854775807, int step=1) -> Tensor(a)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto in = args[0].ITensor(); + auto in = args[0].ITensorOrFreeze(ctx); auto axis = args[1].unwrapToInt(); auto maxDim = static_cast(in->getDimensions().d[axis]); // Handle case when given tensor index is negative