Skip to content

Commit 4addf3d

Browse files
author
Wei Wei
committed
[fx2trt] prod dim is not wrapped as a tuple for single dim
Summary: as titled CI test did not caught this failure before Reviewed By: wushirong Differential Revision: D34932995 fbshipit-source-id: 822ce98fde0545835af3443150275546a946691d
1 parent 0d2cd71 commit 4addf3d

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

fx/converters/converter_utils.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,6 +523,9 @@ def add_reduce_layer(
523523
else:
524524
dim = kwargs["dim"] # type: ignore[assignment]
525525

526+
if not isinstance(dim, Sequence):
527+
dim = (dim,)
528+
526529
if not network.has_implicit_batch_dimension:
527530
dim = tuple(len(input_val.shape) + i if i < 0 else i for i in dim)
528531
else:

0 commit comments

Comments
 (0)