update code

alibaba · Jan 14, 2025 · b8f5a04 · b8f5a04
1 parent d1bd022
commit b8f5a04
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/tzrec/acc/trt_utils.py b/tzrec/acc/trt_utils.py
@@ -197,8 +197,6 @@ def export_model_trt(
         if v.size(0) < 2:
             v = torch.zeros((2,) + v.size()[1:], device="cuda:0", dtype=v.dtype)
         values_list_cuda.append(v)
-        print(v.is_contiguous(memory_format=torch.contiguous_format))
-        print(v.is_contiguous(memory_format=torch.channels_last))
         dynamic_shapes_list.append(dict_dy)
 
     # convert dense

diff --git a/tzrec/tests/rank_integration_test.py b/tzrec/tests/rank_integration_test.py
@@ -11,6 +11,7 @@
 
 import json
 import os
+import shutil
 import tempfile
 import unittest
 
@@ -33,9 +34,9 @@ def setUp(self):
         os.chmod(self.test_dir, 0o755)
 
     def tearDown(self):
-        # if self.success:
-        #     if os.path.exists(self.test_dir):
-        #         shutil.rmtree(self.test_dir)
+        if self.success:
+            if os.path.exists(self.test_dir):
+                shutil.rmtree(self.test_dir)
         os.environ.pop("QUANT_EMB", None)
         os.environ.pop("INPUT_TILE", None)
         os.environ.pop("ENABLE_TRT", None)