fix aot ut error, skip on cpu

alibaba · Jan 7, 2025 · 7fd3a37 · 7fd3a37
1 parent ffebbb7
commit 7fd3a37
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tzrec/tests/rank_integration_test.py b/tzrec/tests/rank_integration_test.py
@@ -68,6 +68,7 @@ def _test_rank_nofg(self, pipeline_config_path, reserved_columns, output_columns
             os.path.exists(os.path.join(self.test_dir, "export/scripted_model.pt"))
         )
 
+    @unittest.skipIf(not torch.cuda.is_available(), "cuda not found")
     def test_aot_export(self):
         pipeline_config_path = "tzrec/tests/configs/multi_tower_din_mock.config"
         self.success = utils.test_train_eval(pipeline_config_path, self.test_dir)