fix ut

xadupre · xadupre · commit 5003335097f5 · 2025-05-05T22:51:38.000+02:00
diff --git a/_unittests/ut_torch_export_patches/test_patch_module.py b/_unittests/ut_torch_export_patches/test_patch_module.py
@@ -355,7 +355,7 @@ def test_rewrite_PLBartEncoderLayer(self):
         rewritten = transform_method(PLBartEncoderLayer.forward, verbose=self.verbose)
         self.assertIn(
             (
-                "hidden_states, = torch.cond(hidden_states.dtype == torch.float16 and "
+                "torch.cond(hidden_states.dtype == torch.float16 and "
                 "(torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()), "
                 "branch_cond_then_1, branch_cond_else_1, [hidden_states])"
             ),

Original file line number	Diff line number	Diff line change
`@@ -355,7 +355,7 @@ def test_rewrite_PLBartEncoderLayer(self):`
`355`	`355`	`rewritten = transform_method(PLBartEncoderLayer.forward, verbose=self.verbose)`
`356`	`356`	`self.assertIn(`
`357`	`357`	`(`
`358`		`- "hidden_states, = torch.cond(hidden_states.dtype == torch.float16 and "`
	`358`	`+ "torch.cond(hidden_states.dtype == torch.float16 and "`
`359`	`359`	`"(torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()), "`
`360`	`360`	`"branch_cond_then_1, branch_cond_else_1, [hidden_states])"`
`361`	`361`	`),`