test_gpu.sh

source /opt/dtk/env.sh
export LD_LIBRARY_PATH=/usr/local/lib/python3.7/site-packages/torch/lib/:$LD_LIBRARY_PATH

#export HIP_VISIBLE_DEVICES=0,1,2,3 
DATA=./cifar_dataset pytest  tests
#DATA=./cifar_dataset pytest -v tests/test_zero_tensor_parallel/