[Relax][PyTorch] Fix KeyError: dtype when converting PyTorch model with gradient checkpointing using torch.export (#18461)

This PR is trying to fix issues
#18439.

Co-authored-by: cchung100m <[email protected]>

Author: cchung100m

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -2036,7 +2036,11 @@ def _arange(self, node: fx.Node) -> relax.Var:
         return self.block_builder.emit(relax.op.arange(*start_end_step, dtype=dtype))
 
     def _empty(self, node: fx.Node) -> relax.Var:
-        dtype = self._convert_data_type(str(node.kwargs["dtype"]), self.env)
+        import torch
+
+        dtype = self._convert_data_type(
+            node.kwargs.get("dtype", torch.get_default_dtype()), self.env
+        )
         return self.block_builder.emit(relax.op.zeros(node.args[0], dtype))
 
     def _empty_like(self, node: fx.Node) -> relax.Var:

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -1143,6 +1143,7 @@ def create_convert_map(
             "_assert_tensor_metadata.default": lambda node: self.env[
                 node.args[0]
             ],  # metadata assertion: no-op
+            "empty.default": self._empty,
             "empty.memory_format": self._empty,
             "empty_permuted.default": self._empty,  # Similar to empty with permuted layout
             "empty_like.default": self._empty_like,

tests/python/relax/test_frontend_from_exported_program.py

@@ -5278,6 +5278,27 @@ def main(
     verify_model(Empty(), example_args, {}, Expected, run_ep_decomposition=True)
 
 
+def test_empty_without_dtype():
+    class EmptyWithoutDtype(Module):
+        def forward(self, input):
+            return torch.empty((5, 5))
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(
+            input: R.Tensor((10, 10), dtype="float32")
+        ) -> R.Tuple(R.Tensor((5, 5), dtype="float32")):
+            with R.dataflow():
+                lv: R.Tensor((5, 5), dtype="float32") = R.zeros(R.shape([5, 5]), dtype="float32")
+                gv: R.Tuple(R.Tensor((5, 5), dtype="float32")) = (lv,)
+                R.output(gv)
+            return gv
+
+    example_args = (torch.randn(10, 10, dtype=torch.float32),)
+    verify_model(EmptyWithoutDtype(), example_args, {}, Expected, run_ep_decomposition=True)
+
+
 def test_fill():
     class Fill(Module):
         def forward(self, input: torch.Tensor):