[bugfix] [Relay] fix broadcast in PyTorch frontend (apache#14885)

* fix broadcast_tensors * Update pytorch.py * Update test_forward.py * Update test_forward.py
mei-ye · Jun 1, 2023 · 46fc528 · 46fc528
1 parent e0bdd13
commit 46fc528
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -2260,7 +2260,12 @@ def broadcast_tensors(self, inputs, input_types):
         tensor_list = inputs[0]
         import torch
 
-        res_shape = list(torch.broadcast_shapes(*[self.infer_shape(t) for t in tensor_list]))
+        infer_shape_value = [self.infer_shape(t) for t in tensor_list]
+        # "torch.broadcast_shapes" is available after PyTorch 1.8.0
+        if hasattr(torch, "broadcast_shapes"):
+            res_shape = list(torch.broadcast_shapes(*infer_shape_value))
+        else:
+            res_shape = list(torch.broadcast_tensors(*map(torch.empty, infer_shape_value))[0].shape)
         return [_op.broadcast_to(tensor, res_shape) for tensor in tensor_list]
 
     def Bool(self, inputs, input_types):