From 4bbbef3523dec78f30069ebe4f284cf76a874674 Mon Sep 17 00:00:00 2001 From: Michael Feliz Date: Wed, 5 Oct 2022 08:30:00 -0700 Subject: [PATCH] Fix incorrect casting behavior in floor_divide (#57) # Description Remove incorrect cast to float in the floor_divide converter that created type-mismatches between the converted and original models. Fixes # (https://jira.robot.car/browse/MLA-6918) ## Type of change Please delete options that are not relevant and/or add your own. - Bug fix (non-breaking change which fixes an issue) - New feature (non-breaking change which adds functionality) - Breaking change (fix or feature that would cause existing functionality to not work as expected) - This change requires a documentation update # Checklist: - [ ] My code follows the style guidelines of this project (You can use the linters) - [ ] I have performed a self-review of my own code - [ ] I have commented my code, particularly in hard-to-understand areas and hacks - [ ] I have made corresponding changes to the documentation - [ ] I have added tests to verify my fix or my feature - [ ] New and existing unit tests pass locally with my changes - [ ] I have added the relevant labels to my PR in so that relevant reviewers are notified --- core/conversion/converters/impl/element_wise.cpp | 3 +-- tests/core/conversion/converters/test_element_wise.cpp | 1 + 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/core/conversion/converters/impl/element_wise.cpp b/core/conversion/converters/impl/element_wise.cpp index 274cbe6a70..fd11246855 100644 --- a/core/conversion/converters/impl/element_wise.cpp +++ b/core/conversion/converters/impl/element_wise.cpp @@ -566,8 +566,7 @@ auto element_wise_registrations TORCHTRT_UNUSED = [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { // TODO: Remove with functionalization auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); auto floor_divide = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); diff --git a/tests/core/conversion/converters/test_element_wise.cpp b/tests/core/conversion/converters/test_element_wise.cpp index 712b9ade11..c4b170ed42 100644 --- a/tests/core/conversion/converters/test_element_wise.cpp +++ b/tests/core/conversion/converters/test_element_wise.cpp @@ -298,6 +298,7 @@ TEST(Converters, ATenFloorDivideWithScalarConvertsCorrectly) { %1 : Tensor = aten::floor_divide(%0, %scalar) return (%1))IR"; pointwise_test_helper(graph, true); + pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt); } TEST(Converters, ATenMaxConvertsCorrectly) {