py-horovod: fix compilation with mxnet enabled (#23369)

This commit is contained in:
Andrew W Elble 2021-04-30 19:55:44 -04:00 committed by GitHub
parent a618594c28
commit bdfdba7d73
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -171,6 +171,8 @@ def setup_build_environment(self, env):
env.set('HOROVOD_WITHOUT_PYTORCH', 1) env.set('HOROVOD_WITHOUT_PYTORCH', 1)
if 'frameworks=mxnet' in self.spec: if 'frameworks=mxnet' in self.spec:
env.set('HOROVOD_WITH_MXNET', 1) env.set('HOROVOD_WITH_MXNET', 1)
env.set('MXNET_INCLUDE_PATH', self.spec['mxnet'].prefix.include)
env.set('MXNET_LIBRARY_PATH', join_path(self.spec['mxnet'].libs[0]))
else: else:
env.set('HOROVOD_WITHOUT_MXNET', 1) env.set('HOROVOD_WITHOUT_MXNET', 1)