(tvm) branch main updated: [Relax][PyTorch][Fix] use`_convert_torch_tensor_to_relax()` where possible (#17335)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new e65aab6a4f [Relax][PyTorch][Fix] use`_convert_torch_tensor_to_relax()` 
where possible (#17335)
e65aab6a4f is described below

commit e65aab6a4f55f4b405ef2713f842d6a3b761151b
Author: Masahiro Hiramori <mhg00...@gmail.com>
AuthorDate: Thu Sep 5 22:30:12 2024 +0900

    [Relax][PyTorch][Fix] use`_convert_torch_tensor_to_relax()` where possible 
(#17335)
    
    * use `_convert_torch_tensor_to_relax` where possible
    
    * add type annotation
---
 python/tvm/relax/frontend/torch/fx_translator.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/python/tvm/relax/frontend/torch/fx_translator.py 
b/python/tvm/relax/frontend/torch/fx_translator.py
index 21a0b2d564..6e60c3bb6f 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -62,7 +62,7 @@ class TorchFXImporter:
         return attr_itr
 
     @staticmethod
-    def _convert_data_type(input_type, env: Optional[Dict] = None):
+    def _convert_data_type(input_type: Union[str, torch.dtype], env: 
Optional[Dict] = None):
         """converts the PyTorch scalar type input_type to a TVM dtype."""
         import torch  # type: ignore
 
@@ -1206,9 +1206,8 @@ class TorchFXImporter:
         module = self.named_modules[node.target]
         weight = self.params[module.weight]
         bias = self.params[module.bias]
-        dtype = 
TorchFXImporter._convert_data_type(str(module.running_mean.dtype))
-        running_mean = relax.const(module.running_mean.cpu().detach().numpy(), 
dtype)
-        running_var = relax.const(module.running_var.cpu().detach().numpy(), 
dtype)
+        running_mean = self._convert_torch_tensor_to_relax(module.running_mean)
+        running_var = self._convert_torch_tensor_to_relax(module.running_var)
         eps = module.eps
 
         res_tuple = self.block_builder.emit(
@@ -1769,7 +1768,7 @@ class TorchFXImporter:
                     dtype = self._convert_data_type(str(param.data.dtype))
                     if dtype in ("float32", "float16"):
                         if not keep_params_as_input:
-                            self.params[param] = 
relax.const(param.data.cpu().numpy(), dtype)
+                            self.params[param] = 
self._convert_torch_tensor_to_relax(param)
                     else:
                         raise ValueError("Unsupported data type for model 
parameters: %s" % dtype)
                 # Translate the model.