add l2norm source for FusedLAMB

a3ffb8a7 · Kexin Yu · 04927b3a · a3ffb8a7
Commit a3ffb8a7 authored Mar 23, 2020 by Kexin Yu
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

setup.py setup.py +2 -1

No files found.
--- a/setup.py
+++ b/setup.py
@@ -211,7 +211,8 @@ if "--deprecated_fused_lamb" in sys.argv:
        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',