diff --git a/core/conversion/converters/impl/einsum.cpp b/core/conversion/converters/impl/einsum.cpp index fb031f6c38..fdaafe4e33 100644 --- a/core/conversion/converters/impl/einsum.cpp +++ b/core/conversion/converters/impl/einsum.cpp @@ -18,6 +18,13 @@ auto einsum_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pat auto equation = args[0].unwrapToString(); auto in = args[1].IValue()->toListRef(); + TORCHTRT_CHECK( + in.size() <= 2, + "TensorRT currently supports up to 2 input tensors " + << "to einsum but operation had " << in.size() + << " input tensors, please specify torch_executed_ops=[\"aten::einsum\"] " + << "at compilation time to avoid this error."); + std::vector tensors; // Populate vector of ITensor pointers