Updated test_gpu_memory.py

Author: vfdev-5

tests/ignite/contrib/metrics/test_gpu_memory.py

@@ -1,3 +1,5 @@
+import sys
+
 import torch
 
 from ignite.contrib.metrics import GpuMemory
@@ -7,11 +9,7 @@
 
 @pytest.fixture
 def no_site_packages():
-    try:
-        import pynvml
-    except ImportError:
-        yield "no_site_packages"
-        return
+    import pynvml
     import sys
     assert 'pynvml' in sys.modules
     pynvml_module = sys.modules['pynvml']
@@ -23,20 +21,22 @@ def no_site_packages():
     sys.modules['pynvml'] = pynvml_module
 
 
+@pytest.mark.skipif(sys.version[0] == "2", reason="No pynvml for python 2.7")
 def test_no_pynvml_package(no_site_packages):
 
     with pytest.raises(RuntimeError, match="This contrib module requires pynvml to be installed."):
         GpuMemory()
 
 
-@pytest.mark.skipif(torch.cuda.is_available(), reason="Skip if has GPU")
+@pytest.mark.skipif(sys.version[0] == "2" or torch.cuda.is_available(), reason="No pynvml for python 2.7")
 def test_no_gpu():
 
     with pytest.raises(RuntimeError, match="This contrib module requires available GPU"):
         GpuMemory()
 
 
-@pytest.mark.skipif(torch.cuda.is_available(), reason="Skip if has GPU")
+@pytest.mark.skipif(sys.version[0] == "2" or not (torch.cuda.is_available()),
+                    reason="No pynvml for python 2.7 and no GPU")
 def test_gpu_mem_consumption():
 
     gpu_mem = GpuMemory()