Skip to content

Commit 4bbbef3

Browse files
committed
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
1 parent e608bc7 commit 4bbbef3

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

core/conversion/converters/impl/element_wise.cpp

+1-2
Original file line numberDiff line numberDiff line change
@@ -566,8 +566,7 @@ auto element_wise_registrations TORCHTRT_UNUSED =
566566
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
567567
// TODO: Remove with functionalization
568568
auto self = args[0].ITensorOrFreeze(ctx);
569-
auto otherScalar = args[1].unwrapToScalar().to<float>();
570-
auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
569+
auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
571570
auto floor_divide =
572571
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n));
573572
TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n);

tests/core/conversion/converters/test_element_wise.cpp

+1
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,7 @@ TEST(Converters, ATenFloorDivideWithScalarConvertsCorrectly) {
298298
%1 : Tensor = aten::floor_divide(%0, %scalar)
299299
return (%1))IR";
300300
pointwise_test_helper(graph, true);
301+
pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt);
301302
}
302303

303304
TEST(Converters, ATenMaxConvertsCorrectly) {

0 commit comments

Comments
 (0)