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) {