Skip to content

Commit 71a9db5

Browse files
authored
Merge pull request #1392 from mfeliz-cruise/michael.feliz/upstream_floor_divide_cast
[fix] incorrect casting behavior in floor_divide
2 parents c3c4b1c + 4bbbef3 commit 71a9db5

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)