gistfile1.txt

def test_stream_enter_exit(self):
        def fn(x, y):
            s2 = torch.cuda.Stream()
            s1 = torch.cuda.Stream()
            with s2:
                z1 = torch.add(x, y)
            with s1:
                z = torch.add(x, y)
                y = z + 2 + z1

            return y

        inp = (torch.ones(2, 2) + 1, torch.ones(2, 2))
        expected = fn(*inp)
        fn_opt = torch.compile(fn, fullgraph=True)
        actual = fn_opt(*inp)
        self.assertEqual(expected, actual)