File tree Expand file tree Collapse file tree 1 file changed +6
-7
lines changed Expand file tree Collapse file tree 1 file changed +6
-7
lines changed Original file line number Diff line number Diff line change @@ -175,7 +175,9 @@ def get_extensions():
175
175
176
176
define_macros = []
177
177
178
- extra_compile_args = {}
178
+ extra_compile_args = {
179
+ 'cxx' : []
180
+ }
179
181
if (torch .cuda .is_available () and ((CUDA_HOME is not None ) or is_rocm_pytorch )) \
180
182
or os .getenv ('FORCE_CUDA' , '0' ) == '1' :
181
183
extension = CUDAExtension
@@ -190,16 +192,13 @@ def get_extensions():
190
192
else :
191
193
define_macros += [('WITH_HIP' , None )]
192
194
nvcc_flags = []
193
- extra_compile_args = {
194
- 'cxx' : [],
195
- 'nvcc' : nvcc_flags ,
196
- }
195
+ extra_compile_args ['nvcc' ] = nvcc_flags
197
196
198
197
if sys .platform == 'win32' :
199
198
define_macros += [('torchvision_EXPORTS' , None )]
200
-
201
- extra_compile_args .setdefault ('cxx' , [])
202
199
extra_compile_args ['cxx' ].append ('/MP' )
200
+ elif sys .platform == 'linux' :
201
+ extra_compile_args ['cxx' ].append ('-fopenmp' )
203
202
204
203
debug_mode = os .getenv ('DEBUG' , '0' ) == '1'
205
204
if debug_mode :
You can’t perform that action at this time.
0 commit comments