Restore torch defaults between sgl-kernel tests (#11131)

This commit is contained in:
Ben Barsdell
2025-11-04 18:51:23 +11:00
committed by GitHub
parent 173e0f704f
commit fef3a6b63b
2 changed files with 30 additions and 0 deletions

View File

@@ -0,0 +1,14 @@
import pytest
import torch
# This ensures the torch defaults don't get left in modified states between
# tests (e.g., when a test fails before restoring the original value), which
# can cause subsequent tests to fail.
@pytest.fixture(autouse=True)
def reset_torch_defaults():
orig_default_device = torch.get_default_device()
orig_default_dtype = torch.get_default_dtype()
yield
torch.set_default_dtype(orig_default_dtype)
torch.set_default_device(orig_default_device)

View File

@@ -0,0 +1,16 @@
import pytest
import torch
def test_change_torch_defaults():
torch.set_default_device("cpu:0")
torch.set_default_dtype(torch.float16)
def test_check_torch_defaults():
assert torch.get_default_device() == torch.device("cpu")
assert torch.get_default_dtype() == torch.float32
if __name__ == "__main__":
pytest.main([__file__])