Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[LAI] Make lite tests safe for combined package (#15204) 

Make lite tests safe for combined package

Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com>
This commit is contained in:
otaj 2022-10-20 11:10:39 +02:00 committed by GitHub
parent 3791dda609
commit 741462f373
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tests/tests_lite/accelerators/test_cuda.py
View File

@ -11,14 +11,15 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import importlib
import os import os
import sys
from unittest import mock from unittest import mock
import pytest import pytest
import torch import torch
from tests_lite.helpers.runif import RunIf from tests_lite.helpers.runif import RunIf
import lightning_lite
from lightning_lite.accelerators.cuda import CUDAAccelerator, is_cuda_available, num_cuda_devices from lightning_lite.accelerators.cuda import CUDAAccelerator, is_cuda_available, num_cuda_devices
@ -67,9 +68,8 @@ def test_num_cuda_devices_without_nvml(*_):
@mock.patch.dict(os.environ, {}, clear=True) @mock.patch.dict(os.environ, {}, clear=True)
def test_force_nvml_based_cuda_check(monkeypatch): def test_force_nvml_based_cuda_check():
"""Test that we force PyTorch to use the NVML-based CUDA checks.""" """Test that we force PyTorch to use the NVML-based CUDA checks."""
monkeypatch.delitem(sys.modules, "lightning_lite") # 'un-import' lightning_lite importlib.reload(lightning_lite) # reevaluate top-level code, without becoming a different object
import lightning_lite # noqa: F401
assert os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] == "1" assert os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] == "1"