add l2norm source for FusedLAMB

lmw0320 · Mar 23, 2020 · a3ffb8a · a3ffb8a
1 parent 04927b3
commit a3ffb8a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -211,7 +211,8 @@ def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
         ext_modules.append(
             CUDAExtension(name='fused_lamb_cuda',
                           sources=['apex/contrib/csrc/optimizers/fused_lamb_cuda.cpp',
-                                   'apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu'],
+                                   'apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu',
+                                   'csrc/multi_tensor_l2norm_kernel.cu'],
                           include_dirs=[os.path.join(this_dir, 'csrc')],
                           extra_compile_args={'cxx': ['-O3',] + version_dependent_macros,
                                               'nvcc':['-O3',