refine chore code

bytedance · Jul 1, 2024 · a2419af · a2419af
1 parent a15a76e
commit a2419af
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 37 deletions.
diff --git a/runtime/python/brt/utils.py b/runtime/python/brt/utils.py
@@ -16,6 +16,8 @@ def brt_dtype_to_torch_dtype(dtype: brt.DType):
         return torch.float64
     if dtype == brt.DType.bool:
         return torch.bool
+    if dtype == brt.DType.int1:
+        return torch.bool
     if dtype == brt.DType.int8:
         return torch.int8
     if dtype == brt.DType.int16:

diff --git a/tests/compatibilty_test/execute.py b/tests/compatibilty_test/execute.py
@@ -13,14 +13,14 @@
 # ==============================================================================
 
 import brt
+from brt.utils import brt_dtype_to_torch_dtype
 
 import torch
 import numpy as np
 import os
 import re
 
 from reporting import TestResult
-from utils import (mlir_type_to_torch_dtype)
 
 
 class BRTBackend:
@@ -47,7 +47,7 @@ def _generate_torch_outputs(self):
         for offset in self.session.get_output_arg_offsets():
             outputs.append(
                 torch.empty(self.session.get_static_shape(offset),
-                            dtype=mlir_type_to_torch_dtype(
+                            dtype=brt_dtype_to_torch_dtype(
                                 self.session.get_data_type(offset)),
                             device=self.device))
         return outputs

diff --git a/tests/compatibilty_test/utils.py b/tests/compatibilty_test/utils.py