File tree Expand file tree Collapse file tree 1 file changed +12
-0
lines changed Expand file tree Collapse file tree 1 file changed +12
-0
lines changed Original file line number Diff line number Diff line change @@ -46,8 +46,20 @@ if ! [ -x "$(command -v nvidia-smi)" ]; then
46
46
EXTRA_ARGS=" -n auto"
47
47
fi
48
48
49
+ # Lack of HorovodJoin CPU kernels when install Horovod with NCCL
50
+ if [ " $( uname) " != " Darwin" ]; then
51
+ # Mac only with MPI
52
+ python -m pip uninstall horovod -y
53
+ bash /install/install_horovod.sh $HOROVOD_VERSION --only-cpu
54
+ fi
49
55
# TODO(jamesrong): Test on GPU.
50
56
CUDA_VISIBLE_DEVICES=" " mpirun -np 2 -H localhost:2 --allow-run-as-root pytest -v ./tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
57
+ # Reinstall Horovod after tests
58
+ if [ " $( uname) " != " Darwin" ]; then
59
+ # Mac only with MPI
60
+ python -m pip uninstall horovod -y
61
+ bash /install/install_horovod.sh $HOROVOD_VERSION
62
+ fi
51
63
52
64
# Only use GPU 0 if available.
53
65
if [ -x " $( command -v nvidia-smi) " ]; then
You can’t perform that action at this time.
0 commit comments