def deterministic_torch_context():
    """Context manager for deterministic torch settings."""
    # Save previous settings
    prev_deterministic = torch.are_deterministic_algorithms_enabled()

    # Check if fill_uninitialized_memory exists in this PyTorch version
    prev_fill_uninitialized: Optional[bool] = None
    has_fill_uninitialized = hasattr(
        torch.utils.deterministic,
        'fill_uninitialized_memory',
    )

    if has_fill_uninitialized:
        prev_fill_uninitialized = torch.utils.deterministic.fill_uninitialized_memory  # type: ignore
        torch.utils.deterministic.fill_uninitialized_memory = True  # type: ignore

    # Enable deterministic settings
    torch.use_deterministic_algorithms(True)

    try:
        yield
    finally:
        # Restore previous settings
        torch.use_deterministic_algorithms(prev_deterministic)
        if has_fill_uninitialized and prev_fill_uninitialized is not None:
            torch.utils.deterministic.fill_uninitialized_memory = prev_fill_uninitialized  # type: ignore
Context manager for deterministic torch settings.
tests/models/hf/test_hf_fsdp_meta.py