lightning/tests/tests_data/conftest.py

11 lines
333 B
Python

import pytest
import torch.distributed
@pytest.fixture(autouse=True)
def teardown_process_group():
"""Ensures that the distributed process group gets closed before the next test runs."""
yield
if torch.distributed.is_available() and torch.distributed.is_initialized():
torch.distributed.destroy_process_group()