support torch.permute

apache · Jul 22, 2024 · aa2e91a · aa2e91a
1 parent f22958e
commit aa2e91a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -550,7 +550,11 @@ def _flatten(self, node: fx.node.Node) -> relax.Var:
         return self.block_builder.emit(relax.op.reshape(x, new_shape))
 
     def _permute(self, node: fx.node.Node) -> relax.Var:
+        import torch  # type: ignore
+
         args = self.retrieve_args(node)
+        if isinstance(args[1], (torch.Size, tuple, list)):
+            return self.block_builder.emit(relax.op.permute_dims(args[0], tuple(args[1])))
         return self.block_builder.emit(relax.op.permute_dims(args[0], args[1:]))
 
     def _reshape(self, node: fx.node.Node) -> relax.Var: