mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
5cdbcb5233
This reverts commitcde81e92b9. Reverted https://github.com/pytorch/pytorch/pull/164522 on behalf of https://github.com/atalman due to Breaks periodic: test/dynamo/test_streams.py::TestStreams::test_stream_weakref [GH job link](https://github.com/pytorch/pytorch/actions/runs/18909552619/job/53979171605) [HUD commit link](cde81e92b9) ([comment](https://github.com/pytorch/pytorch/pull/164522#issuecomment-3462450571))
41 lines
1.0 KiB
Python
41 lines
1.0 KiB
Python
# Owner(s): ["module: dynamo"]
|
|
import weakref
|
|
|
|
import torch
|
|
import torch._dynamo.test_case
|
|
import torch._dynamo.testing
|
|
from torch.testing._internal.common_utils import requires_cuda
|
|
|
|
|
|
class TestStreams(torch._dynamo.test_case.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super().setUpClass()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
super().tearDownClass()
|
|
|
|
def test_stream_weakref(self):
|
|
s = torch.Stream()
|
|
weakref.ref(s)
|
|
|
|
@requires_cuda
|
|
def test_run_opcheck(self):
|
|
from torch._dynamo.variables.streams import fork_stream, join_stream
|
|
from torch.library import opcheck
|
|
|
|
sample_inputs = [
|
|
(0, torch.device("cuda:0"), 1, torch.device("cuda:1")),
|
|
(2, torch.device("cuda:2"), 3, torch.device("cuda:1")),
|
|
]
|
|
for args in sample_inputs:
|
|
opcheck(fork_stream, args)
|
|
opcheck(join_stream, args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from torch._dynamo.test_case import run_tests
|
|
|
|
run_tests()
|