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

View File

@@ -171,6 +171,8 @@ def setup_build_environment(self, env):
env.set('HOROVOD_WITHOUT_PYTORCH', 1)
if 'frameworks=mxnet' in self.spec:
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:
env.set('HOROVOD_WITHOUT_MXNET', 1)