Enable FloatTensor <-> CUDA HalfTensor checks in test_cuda.py

apaszke · apaszke · commit 2e1f8344f718 · 2017-11-24T04:57:48.000-05:00
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -384,12 +384,18 @@ def get_cycles_per_ms():
     return _cycles_per_ms
 
 
-def compare_cpu_gpu(tensor_constructor, arg_constructor, fn, t, precision=1e-5):
+def compare_cpu_gpu(tensor_constructor, arg_constructor, fn, t, precision=1e-5, force_gpu_half=False):
     def tmp(self):
         cpu_tensor = tensor_constructor(t)
-        gpu_tensor = to_gpu(cpu_tensor)
+        type_map = {}
+        if force_gpu_half:
+            type_map = {
+                'torch.FloatTensor': 'torch.cuda.HalfTensor',
+                'torch.DoubleTensor': 'torch.cuda.HalfTensor',
+            }
+        gpu_tensor = to_gpu(cpu_tensor, type_map)
         cpu_args = arg_constructor(t)
-        gpu_args = [to_gpu(arg) for arg in cpu_args]
+        gpu_args = [to_gpu(arg, type_map) for arg in cpu_args]
         cpu_result = getattr(cpu_tensor, fn)(*cpu_args)
         try:
             gpu_result = getattr(gpu_tensor, fn)(*gpu_args)
@@ -1100,6 +1106,8 @@ def test_nvtx(self):
 
                 assert not hasattr(TestCuda, test_name), "Duplicated test name: " + test_name
                 setattr(TestCuda, test_name, compare_cpu_gpu(constr, arg_constr, name_inner, t, precision))
+                if t == torch.FloatTensor:
+                    setattr(TestCuda, test_name + '_gpu_half', compare_cpu_gpu(constr, arg_constr, name_inner, t, precision, force_gpu_half=True))
 
 
 if __name__ == '__main__':
