Merge pull request #138 from lcskrishna/cl/enable-nn-test

enable test_nn unit tests.

Author: iotamudelta

test/common_cuda.py

@@ -2,12 +2,13 @@
 
 import torch
 import torch.cuda
+from common import TEST_WITH_ROCM
 
 
 TEST_CUDA = torch.cuda.is_available()
 TEST_MULTIGPU = TEST_CUDA and torch.cuda.device_count() >= 2
 CUDA_DEVICE = TEST_CUDA and torch.device("cuda:0")
-TEST_CUDNN = TEST_CUDA and torch.backends.cudnn.is_acceptable(torch.tensor(1., device=CUDA_DEVICE))
+TEST_CUDNN = TEST_CUDA and (TEST_WITH_ROCM or torch.backends.cudnn.is_acceptable(torch.tensor(1., device=CUDA_DEVICE)))
 TEST_CUDNN_VERSION = TEST_CUDNN and torch.backends.cudnn.version()
 
 

test/common_nn.py

@@ -40,7 +40,8 @@ def get_weight(m):
         module_name='Linear',
         constructor_args=(10, 8),
         input_size=(4, 10),
-        reference_fn=lambda i, p: torch.mm(i, p[0].t()) + p[1].view(1, -1).expand(4, 8)
+        reference_fn=lambda i, p: torch.mm(i, p[0].t()) + p[1].view(1, -1).expand(4, 8),
+        test_cuda = (not TEST_WITH_ROCM)
     ),
     dict(
         module_name='Linear',
@@ -115,6 +116,7 @@ def get_weight(m):
         constructor_args=(1,),
         input_size=(10, 20),
         reference_fn=lambda i, _: torch.exp(i).div_(torch.exp(i).sum(1, True).expand(10, 20)).log_(),
+        test_cuda = (not TEST_WITH_ROCM)
     ),
     dict(
         module_name='LogSoftmax',
@@ -128,7 +130,8 @@ def get_weight(m):
         module_name='ELU',
         constructor_args=(2.,),
         input_size=(3, 2, 5),
-        reference_fn=lambda x, _: torch.where(x >= 0, x, 2 * (x.exp() - 1))
+        reference_fn=lambda x, _: torch.where(x >= 0, x, 2 * (x.exp() - 1)),
+        test_cuda=(not TEST_WITH_ROCM),
     ),
     # TODO: reference function
     dict(
@@ -254,7 +257,8 @@ def get_weight(m):
     ),
     dict(
         module_name='Tanhshrink',
-        input_size=(2, 3, 4, 5)
+        input_size=(2, 3, 4, 5),
+        test_cuda = (not TEST_WITH_ROCM)
     ),
 ]
 

test/run_test.py

@@ -47,7 +47,6 @@
     'distributions',
     'multiprocessing',
     'nccl',
-    'nn',
     'utils',
 ]