Skip to content

Commit f64f4c1

Browse files
Yinghai LuWei Wei
authored andcommitted
Handle None input tensors in lowering if they are eventually not used (#85)
Summary: Pull Request resolved: https://github.com/pytorch/fx2trt/pull/85 Corner case. Reviewed By: khabinov Differential Revision: D36727701 fbshipit-source-id: f2a105d70a1f1a04395ae08e73ae614a080947a2
1 parent cfd5e56 commit f64f4c1

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

fx/lower.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,10 @@ def __call__(
230230
== LowerPrecision.FP16
231231
):
232232
module.half()
233-
inputs = tuple(x.half() if x.dtype == torch.float32 else x for x in inputs)
233+
inputs = tuple(
234+
x.half() if x is not None and x.dtype == torch.float32 else x
235+
for x in inputs
236+
)
234237
pm = self.lower_pass_manager_builder.build_lower_pipeline(inputs)
235238

236239
lower_result = pm(module)

0 commit comments

Comments
 (0)