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