Skip to content

Commit

Permalink
fix(aten::slice): Allow slicing of pytorch tensors
Browse files Browse the repository at this point in the history
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 <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>
  • Loading branch information
narendasan committed Jul 28, 2021
1 parent af36b5b commit 50f012e
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion core/conversion/converters/impl/select.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<int64_t>(in->getDimensions().d[axis]);
// Handle case when given tensor index is negative
Expand Down

0 comments on commit 50f012e

Please sign in to comment.