extend repeat_interleave op for relay.Expr (#8839)

Co-authored-by: Valery Chernov <valery.chernov@deelvin.com>
apache · Aug 27, 2021 · b4b194d · b4b194d
1 parent 227bf7f
commit b4b194d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -573,6 +573,12 @@ def repeat_interleave(self, inputs, input_types):
         if isinstance(inputs[1], int):
             repeats = inputs[1]
             axis = inputs[2]
+        elif isinstance(inputs[1], _expr.Expr):
+            if isinstance(inputs[1], _expr.Constant):
+                repeats = int(inputs[1].data.numpy())
+            else:
+                repeats, _ = try_infer_value(inputs[1], lambda ret: ret.tolist())
+            axis = inputs[2]
         else:
             msg = "Only repeat with one value as repeat is currently supported."
             raise AssertionError(msg)